"""
MCP 客户端桥接
==============
连接外部 MCP Server，把它们的 tools 提供给自动回复链路调用，
让客服 AI 能查文档/挂号/业务系统等，回复更精准。

服务器列表配置在 ai_config.AI_MCP_SERVERS，例如：

[
  {
    "name": "filesystem",
    "enabled": true,
    "transport": "stdio",
    "command": "npx",
    "args": ["-y", "@modelcontextprotocol/server-filesystem", "D:/docs"],
    "env": {},
    "cwd": null
  }
]
"""

from __future__ import annotations

import asyncio
import json
import os
import re
from contextlib import AsyncExitStack
from typing import Any, Optional

import ai_config

# OpenAI / DeepSeek function name 只允许 [a-zA-Z0-9_-]
_SAFE = re.compile(r"[^a-zA-Z0-9_-]+")


def _safe_name(server: str, tool: str) -> str:
    s = _SAFE.sub("_", server.strip()) or "srv"
    t = _SAFE.sub("_", tool.strip()) or "tool"
    return f"{s}__{t}"


def _parse_name(namespaced: str) -> tuple[str, str]:
    if "__" not in namespaced:
        raise ValueError(f"invalid tool name (expect server__tool): {namespaced}")
    server, tool = namespaced.split("__", 1)
    return server, tool


def _enabled_servers() -> list[dict]:
    raw = getattr(ai_config, "AI_MCP_SERVERS", None) or []
    out = []
    for i, s in enumerate(raw):
        if not isinstance(s, dict):
            continue
        if s.get("enabled", True) is False:
            continue
        name = (s.get("name") or f"server{i}").strip()
        transport = (s.get("transport") or "stdio").strip().lower()
        if transport == "stdio" and not s.get("command"):
            continue
        if transport in ("sse", "http", "streamable_http") and not s.get("url"):
            continue
        cfg = dict(s)
        cfg["name"] = name
        cfg["transport"] = transport
        out.append(cfg)
    return out


def _tool_result_to_text(result: Any) -> str:
    """把 MCP CallToolResult 转成发给模型的纯文本。"""
    try:
        if result is None:
            return ""
        # CallToolResult: content list of TextContent / ImageContent / ...
        parts = []
        content = getattr(result, "content", None)
        if content is None and isinstance(result, dict):
            content = result.get("content")
        if not content:
            # 兜底序列化
            return json.dumps(result, ensure_ascii=False, default=str)[:8000]
        for block in content:
            btype = getattr(block, "type", None) or (block.get("type") if isinstance(block, dict) else None)
            if btype == "text":
                text = getattr(block, "text", None) or (block.get("text") if isinstance(block, dict) else "")
                parts.append(text or "")
            else:
                parts.append(json.dumps(block, ensure_ascii=False, default=str))
        text = "\n".join(p for p in parts if p)
        is_err = getattr(result, "isError", False)
        if is_err:
            return f"[tool error]\n{text}"[:8000]
        return text[:8000]
    except Exception as e:
        return f"[tool result parse error] {e}"


class McpHub:
    """
    一次回复周期内的 MCP 连接池。
    用法：
      async with McpHub() as hub:
          tools = hub.openai_tools()
          text = await hub.call_tool("srv__search", {"q": "..."})
    """

    def __init__(self, servers: list[dict] | None = None):
        self._cfgs = servers if servers is not None else _enabled_servers()
        self._stack: Optional[AsyncExitStack] = None
        self._sessions: dict[str, Any] = {}          # server_name -> ClientSession
        self._tool_meta: dict[str, dict] = {}        # namespaced -> {server, tool, schema...}
        self._openai_tools: list[dict] = []

    @property
    def server_names(self) -> list[str]:
        return list(self._sessions.keys())

    @property
    def tool_count(self) -> int:
        return len(self._openai_tools)

    def openai_tools(self) -> list[dict]:
        return list(self._openai_tools)

    async def __aenter__(self) -> "McpHub":
        self._stack = AsyncExitStack()
        await self._stack.__aenter__()
        for cfg in self._cfgs:
            try:
                await self._connect(cfg)
            except Exception as e:
                print(f"    [MCP] 连接失败 [{cfg.get('name')}]: {e}")
        return self

    async def __aexit__(self, *exc):
        if self._stack is not None:
            await self._stack.__aexit__(*exc)
        self._sessions.clear()
        self._tool_meta.clear()
        self._openai_tools.clear()
        self._stack = None

    async def _connect(self, cfg: dict):
        name = cfg["name"]
        transport = cfg["transport"]

        if transport == "stdio":
            from mcp import ClientSession, StdioServerParameters
            from mcp.client.stdio import stdio_client

            env = os.environ.copy()
            extra = cfg.get("env") or {}
            env.update({str(k): str(v) for k, v in extra.items()})
            params = StdioServerParameters(
                command=cfg["command"],
                args=list(cfg.get("args") or []),
                env=env,
                cwd=cfg.get("cwd") or None,
            )
            read, write = await self._stack.enter_async_context(stdio_client(params))
            session = await self._stack.enter_async_context(ClientSession(read, write))
        elif transport in ("sse", "http", "streamable_http"):
            # 优先 streamable HTTP，其次 SSE（视 SDK 版本而定）
            url = cfg["url"]
            headers = cfg.get("headers") or {}
            session = None
            try:
                from mcp import ClientSession
                from mcp.client.streamable_http import streamablehttp_client
                read, write, _ = await self._stack.enter_async_context(
                    streamablehttp_client(url, headers=headers)
                )
                session = await self._stack.enter_async_context(ClientSession(read, write))
            except Exception:
                from mcp import ClientSession
                from mcp.client.sse import sse_client
                read, write = await self._stack.enter_async_context(
                    sse_client(url, headers=headers)
                )
                session = await self._stack.enter_async_context(ClientSession(read, write))
        else:
            raise ValueError(f"unsupported transport: {transport}")

        await session.initialize()
        self._sessions[name] = session

        listed = await session.list_tools()
        tools = getattr(listed, "tools", None) or []
        for tool in tools:
            tname = getattr(tool, "name", "") or ""
            if not tname:
                continue
            ns = _safe_name(name, tname)
            desc = getattr(tool, "description", None) or f"MCP tool {tname} from {name}"
            schema = getattr(tool, "inputSchema", None) or {"type": "object", "properties": {}}
            if isinstance(schema, dict):
                params = schema
            else:
                # pydantic model
                try:
                    params = schema if isinstance(schema, dict) else dict(schema)
                except Exception:
                    params = {"type": "object", "properties": {}}

            self._tool_meta[ns] = {
                "server": name,
                "tool": tname,
                "description": desc,
            }
            self._openai_tools.append({
                "type": "function",
                "function": {
                    "name": ns,
                    "description": f"[{name}] {desc}",
                    "parameters": params,
                },
            })
        print(f"    [MCP] 已连接 [{name}]，工具 {len(tools)} 个")

    async def call_tool(self, namespaced: str, arguments: dict | None = None) -> str:
        server, tool = _parse_name(namespaced)
        # 若命名被 safe 化，用 meta 反查真实 tool 名
        meta = self._tool_meta.get(namespaced)
        if meta:
            server = meta["server"]
            tool = meta["tool"]
        session = self._sessions.get(server)
        if session is None:
            return f"[MCP] server not connected: {server}"
        try:
            result = await session.call_tool(tool, arguments=arguments or {})
            return _tool_result_to_text(result)
        except Exception as e:
            return f"[MCP] call_tool failed: {e}"


def run_coro(coro):
    """在同步代码里跑一段 async（回复线程用）。"""
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = None
    if loop and loop.is_running():
        # 极少见：已在事件循环中 → 新开线程跑
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            return pool.submit(asyncio.run, coro).result()
    return asyncio.run(coro)


async def probe_servers(servers: list[dict] | None = None) -> dict:
    """探测 MCP 服务器连通性与工具列表（供 GUI 测试）。"""
    async with McpHub(servers) as hub:
        tools = [
            {"name": n, "server": m["server"], "tool": m["tool"], "description": m["description"]}
            for n, m in hub._tool_meta.items()
        ]
        return {
            "ok": True,
            "servers": hub.server_names,
            "tool_count": hub.tool_count,
            "tools": tools,
        }