import os
import asyncio
from contextlib import AsyncExitStack
from typing import List, Dict, Any

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamablehttp_client

from .config import MCPConfig, MCPServerConfig
from .logger import log_event
from tools.base import BaseTool
from tools.mcp_tool import MCPTool


class MCPManager:
    def __init__(self, config: MCPConfig):
        self.config = config
        self.exit_stack = AsyncExitStack()
        self.sessions: List[ClientSession] = []
        self.tools: List[BaseTool] = []
        self._ready = False
        self._ready_event = asyncio.Event()

    @property
    def is_ready(self) -> bool:
        return self._ready

    async def wait_until_ready(self, timeout: float = 30.0) -> bool:
        """等待 MCP 连接就绪，返回是否成功"""
        try:
            await asyncio.wait_for(self._ready_event.wait(), timeout=timeout)
            return True
        except asyncio.TimeoutError:
            return False

    async def connect_all(self):
        """
        Connect to all enabled MCP servers defined in config.
        """
        for server_conf in self.config.mcp_servers:
            if not server_conf.enabled:
                continue

            # Dispatch based on type
            try:
                if server_conf.type == "stdio":
                    await self._connect_stdio(server_conf)
                elif server_conf.type in ["streamableHttp", "sse"]:
                    await self._connect_streamable_http(server_conf)
                elif server_conf.type == "http":
                    log_event(
                        "mcp_server_skipped",
                        server=server_conf.name,
                        reason=f"Type '{server_conf.type}' not yet supported",
                    )
                    continue
                else:
                    log_event(
                        "mcp_server_error",
                        server=server_conf.name,
                        error=f"Unknown MCP server type '{server_conf.type}'",
                    )
            except Exception as e:
                log_event(
                    "mcp_server_error",
                    server=server_conf.name,
                    error=str(e),
                )

        self._ready = True
        self._ready_event.set()

    async def _connect_stdio(self, server_conf: MCPServerConfig):
        """
        Connect to a stdio-based MCP server.
        """
        if not server_conf.command:
            raise ValueError(
                f"MCP server '{server_conf.name}' configured as stdio but missing 'command'"
            )

        env = os.environ.copy()
        if server_conf.env:
            env.update(server_conf.env)

        # Type assertion to satisfy static checkers, since we validated above
        cmd = server_conf.command if server_conf.command else ""

        server_params = StdioServerParameters(
            command=cmd, args=server_conf.args, env=env
        )

        read, write = await self.exit_stack.enter_async_context(
            stdio_client(server_params)
        )
        session = await self.exit_stack.enter_async_context(ClientSession(read, write))

        await session.initialize()
        self.sessions.append(session)

        await self._discover_tools(session, server_conf.name)

        log_event("mcp_server_connected", server=server_conf.name)

    async def _connect_streamable_http(self, server_conf: MCPServerConfig):
        """
        Connect to a streamableHttp/SSE-based MCP server.
        """
        if not server_conf.url:
            raise ValueError(
                f"MCP server '{server_conf.name}' configured as {server_conf.type} but missing 'url'"
            )

        # 支持 headers（用于 API key 等认证）
        headers = server_conf.headers if server_conf.headers else None

        read, write, _ = await self.exit_stack.enter_async_context(
            streamablehttp_client(server_conf.url, headers=headers)
        )
        session = await self.exit_stack.enter_async_context(ClientSession(read, write))

        await session.initialize()
        self.sessions.append(session)

        await self._discover_tools(session, server_conf.name)

        log_event("mcp_server_connected", server=server_conf.name)

    async def _discover_tools(self, session: ClientSession, server_name: str):
        """
        List tools from the session and create MCPTool instances.
        """
        result = await session.list_tools()

        for tool_model in result.tools:
            mcp_tool = MCPTool(
                name=tool_model.name,
                description=f"[{server_name}] {tool_model.description or ''}",
                input_schema=tool_model.inputSchema,
                session=session,
            )
            self.tools.append(mcp_tool)

    def get_tools(self) -> List[BaseTool]:
        return self.tools

    async def cleanup(self):
        """
        Close all connections.
        """
        await self.exit_stack.aclose()
