# -*- coding: utf-8 -*-
"""
MCP Client Manager 模块
管理多个 MCP Server 连接，支持 stdio 传输（HTTP/SSE 暂未实现）。
工具命名规则：mcp__<server_name>__<tool_name>，避免与内置工具冲突。

# 配置文件格式（~/.abyss/mcp.json）：
{
  "mcp_servers": [
    {"name": "github", "transport": "stdio", "command": "npx", "args": [...], "env": {...}}
  ]
}

SDK 桥接：mcp 官方 SDK 全异步，本模块在内部维护后台事件循环线程，
对外暴露同步接口（start_all/call_tool/stop_all），让 abyss 主同步代码无感调用。
"""
import asyncio
import json
import os
import threading
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple


class MCPServerType(Enum):
    """MCP Server 传输类型。"""
    STDIO = "stdio"
    HTTP = "http"
    SSE = "sse"


@dataclass
class ServerConfig:
    """单个 MCP Server 的配置。"""
    name: str
    transport: MCPServerType
    command: Optional[str] = None
    args: List[str] = field(default_factory=list)
    env: Dict[str, str] = field(default_factory=dict)
    url: Optional[str] = None
    headers: Dict[str, str] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> Optional["ServerConfig"]:
        """从配置字典构造 ServerConfig，必填字段缺失返回 None"""
        name = data.get("name", "").strip()
        transport_str = data.get("transport", "").strip().lower()
        if not name or not transport_str:
            return None
        try:
            transport = MCPServerType(transport_str)
        except ValueError:
            return None

        if transport == MCPServerType.STDIO:
            command = data.get("command", "").strip()
            if not command:
                return None
            return cls(
                name=name, transport=transport,
                command=command,
                args=data.get("args", []) or [],
                env=data.get("env", {}) or {},
            )
        if transport in (MCPServerType.HTTP, MCPServerType.SSE):
            url = data.get("url", "").strip()
            if not url:
                return None
            return cls(
                name=name, transport=transport,
                url=url,
                headers=data.get("headers", {}) or {},
            )
        return None


class MCPClientManager:
    """管理所有已注册的 MCP Server 连接与工具发现。
    内部用后台事件循环线程桥接异步 mcp SDK，对外暴露同步接口。
    """

    def __init__(self, config_path: Optional[Path] = None):
        self.servers: Dict[str, ServerConfig] = {}
        # 工具全名 -> (server_name, original_tool_name, schema)
        self._tools: Dict[str, Tuple[str, str, Dict[str, Any]]] = {}
        # 运行中的 SDK 会话：server_name -> ClientSession
        self._sessions: Dict[str, Any] = {}
        # 持久化 async context（stdio_client + ClientSession 的 async with 栈）
        self._cm_stacks: Dict[str, AsyncExitStack] = {}
        # 后台事件循环（mcp SDK 全异步，用独立线程跑）
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._loop_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()

        if config_path is not None:
            self._load_config(config_path)

    def _load_config(self, config_path: Path) -> None:
        """从 JSON 配置文件加载 mcp_servers"""
        if not config_path.exists():
            return
        try:
            with open(config_path, "r", encoding="utf-8-sig") as f:
                data = json.load(f)
        except (json.JSONDecodeError, OSError):
            return
        for server_data in data.get("mcp_servers", []) or []:
            cfg = ServerConfig.from_dict(server_data)
            if cfg is not None:
                self.servers[cfg.name] = cfg

    @staticmethod
    def make_tool_name(server_name: str, tool_name: str) -> str:
        """生成 MCP 工具的全名：mcp__<server>__<tool>"""
        return f"mcp__{server_name}__{tool_name}"

    def register_tool(self, server_name: str, tool_schema: Dict[str, Any]) -> None:
        """注册一个已发现的 MCP 工具"""
        original_name = tool_schema.get("name", "")
        if not original_name:
            return
        full_name = self.make_tool_name(server_name, original_name)
        self._tools[full_name] = (server_name, original_name, tool_schema)

    def get_server_for_tool(self, full_name: str) -> Tuple[Optional[str], Optional[str]]:
        """根据工具全名返回 (server_name, original_tool_name)，未注册返回 (None, None)"""
        entry = self._tools.get(full_name)
        if entry is None:
            return (None, None)
        return (entry[0], entry[1])

    def get_openai_schemas(self) -> List[Dict[str, Any]]:
        """生成所有 MCP 工具的 OpenAI function calling 定义"""
        schemas = []
        for full_name, (_, _, tool_schema) in self._tools.items():
            schemas.append({
                "type": "function",
                "function": {
                    "name": full_name,
                    "description": tool_schema.get("description", ""),
                    "parameters": tool_schema.get("inputSchema", {"type": "object", "properties": {}}),
                }
            })
        return schemas

    def list_tool_names(self) -> List[str]:
        """返回所有已注册 MCP 工具的全名"""
        return list(self._tools.keys())

    def start_all(self) -> None:
        """启动后台事件循环，握手所有 stdio MCP Server 并 discover 工具。失败仅告警不抛。"""
        if self._loop is not None:
            return  # 已启动
        self._loop = asyncio.new_event_loop()
        self._loop_thread = threading.Thread(target=self._loop.run_forever, daemon=True)
        self._loop_thread.start()
        # 在后台循环中跑握手 + discover
        future = asyncio.run_coroutine_threadsafe(self._start_all_async(), self._loop)
        try:
            future.result(timeout=30)
        except Exception as e:
            # 握手失败不致命，工具调用时再报错
            pass

    async def _start_all_async(self) -> None:
        """异步握手所有 stdio server：stdio_client + ClientSession + initialize + list_tools"""
        from mcp import ClientSession
        from mcp.client.stdio import stdio_client, StdioServerParameters

        for name, cfg in self.servers.items():
            if cfg.transport != MCPServerType.STDIO:
                continue
            if name in self._sessions:
                continue
            stack = None
            try:
                stack = AsyncExitStack()
                await stack.__aenter__()
                env = {**os.environ, **cfg.env}
                params = StdioServerParameters(
                    command=cfg.command, args=list(cfg.args), env=env
                )
                read, write = await stack.enter_async_context(stdio_client(params))
                session = await stack.enter_async_context(ClientSession(read, write))
                await session.initialize()
                self._sessions[name] = session
                self._cm_stacks[name] = stack
                # discover 工具并注册
                tools_result = await session.list_tools()
                for tool in tools_result.tools:
                    self.register_tool(name, {
                        "name": tool.name,
                        "description": tool.description or "",
                        "inputSchema": tool.inputSchema or {"type": "object", "properties": {}},
                    })
            except Exception:
                # 单个 server 失败不影响其他，但必须关闭已打开的 AsyncExitStack，
                # 否则泄漏的 stdio_client async generator 会在 GC 时跨 Task 关闭，
                # 触发 'Attempted to exit cancel scope in a different task' RuntimeError。
                if stack is not None:
                    try:
                        await stack.aclose()
                    except Exception:
                        pass
                continue

    def call_tool(self, server_name: str, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """同步调用 MCP 工具，通过后台事件循环提交异步任务。"""
        if server_name not in self._sessions:
            return {"success": False, "error": f"MCP server 未启动: {server_name}"}
        if self._loop is None:
            return {"success": False, "error": "MCP 事件循环未启动"}
        future = asyncio.run_coroutine_threadsafe(
            self._call_tool_async(server_name, tool_name, arguments),
            self._loop,
        )
        try:
            return future.result(timeout=30)
        except Exception as e:
            return {"success": False, "error": f"MCP 调用异常: {e}"}

    async def _call_tool_async(self, server_name: str, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """异步调用单个 MCP 工具，把 SDK result 转成 dict。"""
        session = self._sessions[server_name]
        result = await session.call_tool(tool_name, arguments)
        # SDK result.content 是 list[TextContent | ...]，拼成文本
        text_parts = []
        for block in (result.content or []):
            text = getattr(block, "text", None)
            if text:
                text_parts.append(text)
        return {
            "success": not result.isError,
            "stdout": "\n".join(text_parts),
            "stderr": "" if not result.isError else "\n".join(text_parts),
        }

    def stop_all(self) -> None:
        """关闭所有 MCP 会话 + 停止后台事件循环。幂等。"""
        if self._loop is not None:
            future = asyncio.run_coroutine_threadsafe(self._stop_async(), self._loop)
            try:
                future.result(timeout=5)
            except Exception:
                pass
            self._loop.call_soon_threadsafe(self._loop.stop)
            if self._loop_thread:
                self._loop_thread.join(timeout=2)
        self._sessions.clear()
        self._cm_stacks.clear()
        self._loop = None
        self._loop_thread = None

    async def _stop_async(self) -> None:
        """异步关闭所有 async context 栈（stdio_client + ClientSession）。"""
        for stack in self._cm_stacks.values():
            try:
                await stack.aclose()
            except Exception:
                pass


def create_default_manager() -> MCPClientManager:
    """创建指向 ~/.abyss/mcp.json 的默认 MCPClientManager"""
    return MCPClientManager(config_path=Path.home() / ".abyss" / "mcp.json")
