# 中文注释：
# 文件：echobot/agent/mcp/client.py
# 说明：MCP 客户端管理类

"""MCP Client for connecting to MCP servers."""

import os
import sys
from contextlib import asynccontextmanager
from typing import Any

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

from echobot.config.schema import MCPServerConfig


class MCPClient:
    """Manages connections to MCP servers."""

    def __init__(self, servers: dict[str, MCPServerConfig]):
        """
        Initialize MCP client with server configurations.

        Args:
            servers: Dict of server name -> MCPServerConfig
        """
        self.servers = servers
        self._sessions: dict[str, Any] = {}  # server name -> session
        self._tools: dict[str, list[dict[str, Any]]] = {}  # server name -> tools
        self._configs: dict[str, MCPServerConfig] = {}  # server name -> config for reconnecting

    async def connect_all(self) -> None:
        """Connect to all configured MCP servers."""
        import traceback
        for name, config in self.servers.items():
            try:
                await self._connect(name, config)
            except Exception as e:
                print(f"Failed to connect to MCP server '{name}': {e}", file=sys.stderr)
                traceback.print_exc()

    async def _connect(self, name: str, config: MCPServerConfig) -> None:
        """Connect to a single MCP server."""
        if config.type == "stdio":
            await self._connect_stdio(name, config)
        elif config.type == "http":
            await self._connect_http(name, config)
        else:
            raise ValueError(f"Unknown MCP server type: {config.type}")

    async def _connect_stdio(self, name: str, config: MCPServerConfig) -> None:
        """Connect to a stdio-based MCP server."""
        if not config.command or not config.args:
            raise ValueError(f"stdio server '{name}' requires 'command' and 'args'")

        # Merge env with current environment
        env = os.environ.copy()
        if config.env:
            # Fix env keys that were converted by Pydantic (e.g., m_i_n_i_m_a_x__a_p_i__k_e_y -> MINIMAX_API_KEY)
            for key, value in config.env.items():
                # Debug: print the original key
                print(f"[MCP] Original env key: {key}", file=sys.stderr)
                # Pydantic converts "MINIMAX_API_KEY" to "m_i_n_i_m_a_x__a_p_i__k_e_y"
                # Convert back by splitting on "__" and removing single underscores
                parts = key.split("__")
                fixed_parts = []
                for part in parts:
                    fixed_parts.append(part.replace("_", ""))
                fixed_key = "_".join(fixed_parts).upper()
                print(f"[MCP] Fixed env key: {fixed_key}", file=sys.stderr)
                env[fixed_key] = value

        server_params = StdioServerParameters(
            command=config.command,
            args=config.args,
            env=env,
        )

        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                self._sessions[name] = session
                self._configs[name] = config  # Save config for reconnecting
                # Load tools from this server
                self._tools[name] = await self._list_tools_from_session(session)

    async def _connect_http(self, name: str, config: MCPServerConfig) -> None:
        """Connect to an HTTP-based MCP server."""
        if not config.url:
            raise ValueError(f"http server '{name}' requires 'url'")

        # Resolve headers with environment variables
        headers = {}
        if config.headers:
            for key, value in config.headers.items():
                # Resolve ${VAR} patterns
                if value.startswith("${") and value.endswith("}"):
                    env_var = value[2:-1]
                    headers[key] = os.environ.get(env_var, "")
                else:
                    headers[key] = value

        async with streamable_http_client(config.url, headers) as session:
            self._sessions[name] = session
            # Load tools from this server
            self._tools[name] = await self._list_tools_from_session(session)

    async def _list_tools_from_session(self, session: Any) -> list[dict[str, Any]]:
        """List tools from an MCP session."""
        result = await session.list_tools()
        tools = []
        for tool in result.tools:
            tools.append({
                "name": tool.name,
                "description": tool.description or "",
                "inputSchema": tool.inputSchema,
            })
        return tools

    async def disconnect_all(self) -> None:
        """Disconnect from all MCP servers."""
        # Close each session properly
        for session in self._sessions.values():
            try:
                await session.close()
            except Exception:
                pass
        self._sessions.clear()
        self._tools.clear()
        self._configs.clear()

    def get_all_tools(self) -> list[dict[str, Any]]:
        """Get all tools from all connected servers."""
        all_tools = []
        for tools in self._tools.values():
            all_tools.extend(tools)
        return all_tools

    async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> str:
        """Call a tool by name from any connected server."""
        import sys
        import traceback

        # Find which server has this tool
        server_name = None
        for srv_name, tools in self._tools.items():
            for tool in tools:
                if tool["name"] == name:
                    server_name = srv_name
                    break
            if server_name:
                break

        if not server_name:
            return f"Error: Tool '{name}' not found"

        print(f"[MCP] Calling tool: {name} with args: {arguments}", file=sys.stderr)

        # For stdio servers, we need to reconnect before each call
        # because the connection gets closed after each tool call
        config = self._configs.get(server_name)
        if config and config.type == "stdio":
            try:
                # Reconnect
                print(f"[MCP] Reconnecting to stdio server '{server_name}'", file=sys.stderr)
                await self._connect_stdio(server_name, config)
            except Exception as e:
                error_msg = f"[MCP] Reconnect error: {e}\n{traceback.format_exc()}"
                print(error_msg, file=sys.stderr)
                return f"Error reconnecting to MCP server: {e}"

        # Get the session again after reconnect
        session = self._sessions.get(server_name)
        if not session:
            return f"Error: No session for server '{server_name}'"

        try:
            result = await session.call_tool(name, arguments)
            print(f"[MCP] Tool result: {result}", file=sys.stderr)
            # Convert result to string
            if hasattr(result, "content"):
                parts = []
                for content in result.content:
                    if hasattr(content, "text"):
                        parts.append(content.text)
                    else:
                        parts.append(str(content))
                return "\n".join(parts)
            return str(result)
        except Exception as e:
            error_msg = f"[MCP] Tool call error: {e}\n{traceback.format_exc()}"
            print(error_msg, file=sys.stderr)
            return f"Error calling tool '{name}': {e}"

    @property
    def is_connected(self) -> bool:
        """Check if any servers are connected."""
        return len(self._sessions) > 0
