from __future__ import annotations

from typing import Any

from niannian.mcp.client.config import list_enabled_servers, load_mcp_config


class MCPClientPool:
    """对内 MCP Client 池 — lazy connect 外部 MCP。"""

    def __init__(self, config: dict[str, Any] | None = None, *, runtime: str = "auto") -> None:
        self._config = config or load_mcp_config()
        self._runtime = runtime
        self._connections: dict[str, Any] = {}
        mcp_cfg = self._config if "client" in self._config else {}
        global_cfg = mcp_cfg if isinstance(mcp_cfg.get("client"), dict) else {}
        self._max_concurrent = int(global_cfg.get("max_concurrent_servers", 2))
        self._lazy = bool(global_cfg.get("lazy_connect", True))
        self._timeout = int(global_cfg.get("connect_timeout_sec", 30))

    def list_servers(self) -> list[str]:
        return list(list_enabled_servers(self._config, runtime=self._runtime).keys())

    def get_server_config(self, name: str) -> dict[str, Any] | None:
        return list_enabled_servers(self._config, runtime=self._runtime).get(name)

    def list_tools(self, name: str) -> list[str]:
        spec = self.get_server_config(name)
        if spec is None:
            return []
        conn = self._connections.get(name)
        if conn and conn.get("session"):
            return conn["session"].list_tools()
        allow = spec.get("allow_tools")
        if allow:
            return list(allow)
        return []

    def lazy_connect(self, name: str) -> bool:
        if name in self._connections:
            return True
        if name not in self.list_servers():
            return False
        if len(self._connections) >= self._max_concurrent:
            return False
        spec = self.get_server_config(name)
        if spec is None:
            return False
        transport = spec.get("transport", "stdio")
        if transport == "stdio" and spec.get("command"):
            from niannian.mcp.client.stdio_session import StdioMcpSession

            session = StdioMcpSession()
            session.connect(str(spec["command"]), list(spec.get("args") or []))
            self._connections[name] = {"status": "live", "transport": "stdio", "session": session}
            return True
        self._connections[name] = {"status": "stub", "transport": transport}
        return True

    def call_tool(self, name: str, tool: str, args: dict[str, Any]) -> dict[str, Any]:
        spec = self.get_server_config(name)
        if spec is None:
            raise KeyError(f"MCP server not configured: {name}")
        deny = set(spec.get("deny_tools") or [])
        if tool in deny:
            raise PermissionError(f"tool denied: {tool}")
        allow = spec.get("allow_tools")
        if allow and tool not in allow:
            raise PermissionError(f"tool not in allowlist: {tool}")
        if self._lazy and name not in self._connections:
            if not self.lazy_connect(name):
                raise RuntimeError(f"cannot connect to MCP server: {name}")
        conn = self._connections.get(name)
        if conn and conn.get("session"):
            return conn["session"].call_tool(tool, args)
        return {
            "ok": False,
            "stub": True,
            "message": f"external MCP call not spawned in MVP (server={name}, tool={tool})",
            "args": args,
        }

    def disconnect(self, name: str) -> None:
        conn = self._connections.pop(name, None)
        if conn and conn.get("session"):
            conn["session"].close()
