"""MCP Manager - 管理 MCP Server 连接和工具发现"""

from __future__ import annotations

import asyncio
import logging
from typing import Any, TYPE_CHECKING

from agentkit.mcp.client import MCPClient
from agentkit.mcp.transport import HTTPTransport, SSETransport, StdioTransport, Transport
from agentkit.tools.registry import ToolRegistry

if TYPE_CHECKING:
    from agentkit.server.config import MCPServerConfig

logger = logging.getLogger(__name__)


class MCPManager:
    """管理 MCP Server 连接和工具发现

    负责启动/停止 MCP Server 连接，发现远程工具并注册到 ToolRegistry。
    """

    def __init__(
        self,
        configs: dict[str, MCPServerConfig],
        tool_registry: ToolRegistry | None = None,
    ):
        self._configs = configs
        self._tool_registry = tool_registry or ToolRegistry()
        self._clients: dict[str, MCPClient] = {}  # server_name -> MCPClient
        self._transports: dict[str, Transport] = {}  # server_name -> Transport
        self._available: dict[str, bool] = {}  # server_name -> is_available
        self._server_tools: dict[str, list[str]] = {}  # server_name -> [tool_names]

    async def start_all(self) -> None:
        """启动所有配置的 MCP Server，并发发现并注册工具

        使用 asyncio.gather 并发启动，单个服务器失败不影响其他服务器。
        """
        tasks = [
            self._start_server_safe(name, config)
            for name, config in self._configs.items()
        ]
        await asyncio.gather(*tasks)

    async def _start_server_safe(self, name: str, config: MCPServerConfig) -> None:
        """启动单个 MCP Server，失败时标记为不可用"""
        try:
            await self._start_server(name, config)
        except Exception as e:
            logger.error("Failed to start MCP server '%s': %s", name, e)
            self._available[name] = False

    async def _start_server(self, name: str, config: MCPServerConfig) -> None:
        """启动单个 MCP Server"""
        config.validate()

        # 根据配置创建传输层
        if config.transport == "stdio":
            transport = StdioTransport(
                command=config.command,
                args=config.args or [],
                env=config.env,
                timeout=config.timeout,
            )
        elif config.transport == "streamable_http":
            transport = HTTPTransport(
                endpoint=config.url,
                headers=config.headers,
                timeout=config.timeout,
            )
        elif config.transport == "sse":
            transport = SSETransport(
                endpoint=config.url,
                headers=config.headers,
                timeout=config.timeout,
            )
        else:
            raise ValueError(f"Unknown transport: {config.transport}")

        # 建立连接
        await transport.connect()
        self._transports[name] = transport

        # 创建客户端并发现工具
        client = MCPClient.from_transport(transport)
        self._clients[name] = client

        tools = await client.list_tools()
        tool_names = []
        for tool_info in tools:
            tool_name = tool_info.get("name", "")
            tool_desc = tool_info.get("description", "")
            mcp_tool = client.as_tool(tool_name, tool_desc)
            self._tool_registry.register(mcp_tool)
            tool_names.append(tool_name)

        self._server_tools[name] = tool_names
        self._available[name] = True
        logger.info("MCP server '%s' started with tools: %s", name, tool_names)

    async def stop_all(self) -> None:
        """停止所有 MCP Server"""
        for name, transport in self._transports.items():
            try:
                await transport.disconnect()
            except Exception as e:
                logger.error("Error stopping MCP server '%s': %s", name, e)
        self._transports.clear()
        self._clients.clear()
        self._available.clear()
        self._server_tools.clear()

    def is_available(self, server_name: str) -> bool:
        """检查指定 MCP Server 是否可用"""
        return self._available.get(server_name, False)

    def get_server_tools(self, server_name: str) -> list[str]:
        """获取指定 MCP Server 提供的工具列表"""
        return self._server_tools.get(server_name, [])

    def list_all_tools(self) -> list[str]:
        """列出所有 MCP Server 提供的工具"""
        all_tools: list[str] = []
        for tools in self._server_tools.values():
            all_tools.extend(tools)
        return all_tools

    def get_tool_registry(self) -> ToolRegistry:
        """获取工具注册中心"""
        return self._tool_registry
