"""MCP Client - 调用外部 MCP 工具服务器

U16: 重写为 langchain-mcp-adapters 包装层。保留 MCPClient / MCPTool API 向后兼容。
旧的 Transport 注入路径保留（发出 DeprecationWarning），供现有调用方过渡。

传输类型由 URL scheme 自动检测（transport=None 时）：
- ``stdio://command arg1 arg2`` → stdio 传输
- ``http://`` / ``https://``     → streamable_http 传输
- ``sse://...``                  → sse 传输（自动转为 http://）
"""

from __future__ import annotations

import json
import logging
import warnings
from typing import TYPE_CHECKING, Any
from urllib.parse import urlparse

from agentkit.mcp.transport import HTTPTransport, SSETransport, StdioTransport, Transport
from agentkit.tools.base import Tool

if TYPE_CHECKING:
    from langchain_mcp_adapters.client import MultiServerMCPClient

logger = logging.getLogger(__name__)


def _import_langchain_client() -> type["MultiServerMCPClient"]:
    """延迟导入 langchain-mcp-adapters 的 MultiServerMCPClient。

    未安装时抛出带提示信息的 ImportError，便于调用方回退到 Transport 路径。
    """
    try:
        from langchain_mcp_adapters.client import MultiServerMCPClient as _Client
    except ImportError as e:
        raise ImportError(
            "langchain-mcp-adapters 未安装，无法使用 langchain 传输路径。"
            "请执行 pip install 'fischer-agentkit[mcp]' 或 pip install langchain-mcp-adapters。"
            "如需使用旧的 Transport 路径，请通过 MCPClient.from_transport(transport) 创建客户端。"
        ) from e
    return _Client


class MCPClient:
    """MCP Client - 连接外部 MCP Server 并调用工具

    U16: 内部使用 langchain-mcp-adapters 的 MultiServerMCPClient 管理连接。
    保留旧 Transport 注入路径以向后兼容（发出 DeprecationWarning）。

    支持两种模式：
    1. URL scheme 自动检测（推荐）：transport=None，根据 server_url 的 scheme 选择传输
    2. Transport 注入（旧路径，弃用）：传入 Transport 实例，走原有 JSON-RPC 路径
    """

    def __init__(
        self,
        server_url: str,
        timeout: int = 30,
        transport: Transport | None = None,
    ):
        self._server_url = server_url.rstrip("/") if server_url else ""
        self._timeout = timeout
        self._tools_cache: list[dict] | None = None
        self._transport = transport
        # U10 — 懒构造并缓存的 langchain client，避免每次 list_tools/call_tool
        # 都新建 MultiServerMCPClient（stdio 传输下会反复 spawn 子进程）。
        self._lc_client: Any = None

        if transport is not None:
            # 旧 Transport 路径 — 发出 DeprecationWarning，但保持原有行为
            warnings.warn(
                "通过 Transport 实例创建 MCPClient 已弃用，请使用 URL scheme 自动检测"
                "（如 MCPClient('http://...') 或 MCPClient('stdio://...')）。"
                "详见 U16 迁移指南。将在下个版本移除。",
                DeprecationWarning,
                stacklevel=2,
            )
            self._langchain_config: dict[str, Any] | None = None
        else:
            # 新 langchain 路径 — 解析 URL scheme 构建连接配置
            self._langchain_config = self._build_langchain_config(self._server_url, timeout)

    @staticmethod
    def _build_langchain_config(server_url: str, timeout: float) -> dict[str, Any]:
        """根据 URL scheme 构建 langchain-mcp-adapters 连接配置。

        Args:
            server_url: 服务器 URL，支持 stdio:// http:// https:// sse://
            timeout: 超时秒数（仅 http/sse 传输使用）

        Returns:
            langchain-mcp-adapters 连接配置 dict

        Raises:
            ValueError: URL 为空或 scheme 不支持
        """
        if not server_url:
            raise ValueError("server_url 不能为空（transport=None 时需要有效 URL）")

        # stdio 特殊处理：stdio://command arg1 arg2 → command + args
        # 不用 urlparse，因为命令行参数可能含特殊字符
        if server_url.startswith("stdio://"):
            raw = server_url[len("stdio://") :]
            parts = raw.split()
            if not parts:
                raise ValueError(f"无效的 stdio URL: {server_url!r}")
            return {
                "transport": "stdio",
                "command": parts[0],
                "args": parts[1:],
            }

        parsed = urlparse(server_url)
        scheme = parsed.scheme.lower()

        if scheme in ("http", "https"):
            return {
                "transport": "streamable_http",
                "url": server_url,
                "timeout": timeout,
            }

        if scheme == "sse":
            # sse://host/path → http://host/path（sse scheme 不携带 TLS 信息）
            converted = "http://" + server_url[len("sse://") :]
            return {
                "transport": "sse",
                "url": converted,
                "timeout": timeout,
            }

        raise ValueError(
            f"不支持的 URL scheme: {scheme!r}。支持 stdio://, http://, https://, sse://"
        )

    @classmethod
    def from_transport(cls, transport: Transport) -> "MCPClient":
        """从 Transport 实例创建 MCPClient（旧路径，向后兼容）

        发出 DeprecationWarning。建议迁移到 URL scheme 自动检测。
        """
        if isinstance(transport, HTTPTransport):
            server_url = transport.endpoint_url
        elif isinstance(transport, SSETransport):
            server_url = transport.endpoint_url
        elif isinstance(transport, StdioTransport):
            server_url = f"stdio://{transport._command}"
        else:
            server_url = ""
        return cls(server_url=server_url, transport=transport)

    async def _get_lc_client(self) -> Any:
        """懒构造并缓存 langchain ``MultiServerMCPClient`` 实例。

        首次调用时创建，后续返回缓存，避免每次 list_tools/call_tool 都新建
        client（stdio 传输下会 spawn 新子进程，造成连接/进程泄漏）。
        """
        if self._lc_client is None:
            client_cls = _import_langchain_client()
            self._lc_client = client_cls({"server": self._langchain_config})
        return self._lc_client

    async def aclose(self) -> None:
        """关闭缓存的 langchain client（如果它提供 ``aclose`` 方法）。"""
        if self._lc_client is not None:
            aclose = getattr(self._lc_client, "aclose", None)
            if aclose is not None:
                await aclose()
            self._lc_client = None

    async def list_tools(self) -> list[dict]:
        """列出远程 MCP Server 上的工具

        Returns:
            工具列表，每项为 ``{"name":..., "description":..., "inputSchema":...}``

        Raises:
            ImportError: langchain-mcp-adapters 未安装（仅 transport=None 路径）
        """
        if self._transport is not None:
            # 旧 Transport 路径
            if not self._transport.is_connected:
                await self._transport.connect()
            result = await self._transport.send_request("tools/list")
            tools = result.get("tools", []) if isinstance(result, dict) else []
            self._tools_cache = tools
            return self._tools_cache

        # 新 langchain 路径 — 复用缓存的 client
        client = await self._get_lc_client()
        lc_tools = await client.get_tools()
        tools = [
            {
                "name": t.name,
                "description": t.description or "",
                "inputSchema": self._extract_schema(t),
            }
            for t in lc_tools
        ]
        self._tools_cache = tools
        return tools

    @staticmethod
    def _extract_schema(tool: Any) -> dict[str, Any]:
        """从 LangChain Tool 提取 inputSchema（JSON Schema 格式）。

        LangChain 工具通常有 args_schema（pydantic model），回退到 tool.args dict。
        """
        args_schema = getattr(tool, "args_schema", None)
        if args_schema is not None and hasattr(args_schema, "model_json_schema"):
            return args_schema.model_json_schema()
        # ponytail: tool.args 在 langchain BaseTool 上是 dict，够用作回退
        args = getattr(tool, "args", None)
        if isinstance(args, dict):
            return args
        return {"type": "object", "properties": {}}

    async def call_tool(self, tool_name: str, arguments: dict) -> dict:
        """调用远程 MCP 工具

        Args:
            tool_name: 工具名称
            arguments: 工具参数

        Returns:
            MCP 响应格式的 dict：``{"content": [{"type":"text","text":...}]}``

        Raises:
            KeyError: 工具不存在（仅 transport=None 路径）
            ImportError: langchain-mcp-adapters 未安装
        """
        if self._transport is not None:
            # 旧 Transport 路径
            if not self._transport.is_connected:
                await self._transport.connect()
            return await self._transport.send_request(
                "tools/call",
                params={"name": tool_name, "arguments": arguments},
            )

        # 新 langchain 路径 — 复用缓存的 client
        client = await self._get_lc_client()
        lc_tools = await client.get_tools()
        for tool in lc_tools:
            if tool.name == tool_name:
                result = await tool.ainvoke(arguments)
                # 包装为 MCP 响应格式，保持 call_tool 返回形状与旧路径一致
                return {
                    "content": [
                        {
                            "type": "text",
                            "text": json.dumps(result, ensure_ascii=False, default=str),
                        }
                    ]
                }
        raise KeyError(f"MCP 工具不存在: {tool_name!r}")

    def as_tool(self, tool_name: str, description: str = "") -> "MCPTool":
        """将远程 MCP 工具包装为本地 Tool 对象"""
        return MCPTool(
            name=tool_name,
            description=description,
            client=self,
        )


class MCPTool(Tool):
    """MCP 工具 - 通过 MCP Client 调用远程工具"""

    def __init__(
        self,
        name: str,
        description: str,
        client: MCPClient,
        input_schema: dict[str, Any] | None = None,
        output_schema: dict[str, Any] | None = None,
        version: str = "1.0.0",
        tags: list[str] | None = None,
    ):
        super().__init__(
            name=name,
            description=description,
            input_schema=input_schema,
            output_schema=output_schema,
            version=version,
            tags=tags or ["mcp"],
        )
        self._client = client

    async def execute(self, **kwargs) -> dict:
        result = await self._client.call_tool(self.name, kwargs)
        # 解析 MCP 响应格式
        if "content" in result:
            for item in result["content"]:
                if item.get("type") == "text":
                    try:
                        return json.loads(item["text"])
                    except json.JSONDecodeError:
                        return {"result": item["text"]}
        return result