code/mcp/__init__.py

"""MCP (Model Context Protocol) support — full client implementation.

Architecture from Gemini CLI's McpClientManager:
- Manages lifecycle of multiple MCP servers
- Bridges MCP tools into the ToolRegistry as ToolBuilder instances
- Supports stdio and SSE transports
- Lazy initialization (connects on first use)

Based on the MCP Python SDK.
"""

from __future__ import annotations

import asyncio
import json
import logging
import threading
from typing import Any, Callable

from code.config import MCPServerConfig

logger = logging.getLogger(__name__)


# ── MCP Tool Bridge ────────────────────────────────────────────────────


class MCPToolBridge:
    """Wraps an MCP tool as a ToolBuilder, making it indistinguishable from built-in tools.

    This is the key insight from Gemini CLI: MCP tools use the same ToolBuilder
    interface so they get the same policy, confirmation, and scheduling treatment.
    """

    def __init__(
        self,
        server_name: str,
        tool_name: str,
        tool_schema: dict,
        call_fn: Callable,
    ):
        self.server_name = server_name
        self.tool_name = tool_name
        self._schema = tool_schema
        self._call_fn = call_fn

        # Generate MCP-prefixed name
        self.name = f"mcp_{server_name}_{tool_name}"
        self.description = tool_schema.get("description", f"MCP tool: {tool_name}")
        self.kind = "other"
        self.is_read_only = tool_schema.get("annotations", {}).get("readOnlyHint", False)

    def get_tool_definition(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self._schema.get(
                "inputSchema", {"type": "object", "properties": {}}
            ),
            "kind": self.kind,
            "mcp_server": self.server_name,
            "mcp_tool": self.tool_name,
        }

    def call(self, arguments: dict[str, Any]) -> dict[str, Any]:
        """Execute the MCP tool call."""
        return self._call_fn(self.tool_name, arguments)


# ── MCP Client Manager ─────────────────────────────────────────────────


class MCPClientManager:
    """Manages multiple MCP server connections.

    Lifecycle: configure → connect (lazy) → use → disconnect
    """

    def __init__(self):
        self._configs: dict[str, MCPServerConfig] = {}
        self._clients: dict[str, Any] = {}  # server_name → MCP client/session
        self._tools: dict[str, MCPToolBridge] = {}
        self._connected: set[str] = set()
        self._lock = threading.Lock()

    def configure(self, configs: list[MCPServerConfig]) -> None:
        """Set MCP server configurations."""
        self._configs = {c.name: c for c in configs if c.enabled}

    def add_server(self, config: MCPServerConfig) -> None:
        """Add or update a single MCP server config."""
        self._configs[config.name] = config
        # If already connected to this server, disconnect and reconnect
        if config.name in self._connected:
            self._disconnect(config.name)

    def remove_server(self, name: str) -> None:
        """Remove an MCP server."""
        self._configs.pop(name, None)
        self._disconnect(name)

    def connect_all(self) -> dict[str, str | None]:
        """Connect to all configured MCP servers. Returns {name: error_or_None}."""
        results = {}
        for name, config in self._configs.items():
            if config.enabled:
                try:
                    error = self._connect(name, config)
                    results[name] = error
                except Exception as e:
                    results[name] = str(e)
                    logger.error("MCP connect failed for %s: %s", name, e)
        return results

    def disconnect_all(self) -> None:
        """Disconnect all MCP servers."""
        for name in list(self._connected):
            self._disconnect(name)

    def get_tools(self) -> dict[str, MCPToolBridge]:
        """Get all discovered MCP tools."""
        return dict(self._tools)

    def get_tool(self, name: str) -> MCPToolBridge | None:
        """Get a specific MCP tool by its full name (mcp_server_tool)."""
        return self._tools.get(name)

    def call_tool(
        self, server_name: str, tool_name: str, arguments: dict[str, Any]
    ) -> dict[str, Any]:
        """Call an MCP tool on a specific server."""
        client = self._clients.get(server_name)
        if not client:
            # Try lazy connect
            config = self._configs.get(server_name)
            if not config:
                return {"error": f"MCP server '{server_name}' not configured"}
            error = self._connect(server_name, config)
            if error:
                return {"error": f"Failed to connect to '{server_name}': {error}"}
            client = self._clients.get(server_name)
            if not client:
                return {"error": f"MCP server '{server_name}' connection failed"}

        try:
            if hasattr(client, "call_tool"):
                result = client.call_tool(tool_name, arguments)
                if hasattr(result, "content"):
                    # Extract text content from MCP result
                    texts = []
                    for item in result.content:
                        if hasattr(item, "text"):
                            texts.append(item.text)
                        else:
                            texts.append(str(item))
                    return {"output": "\n".join(texts)}
                return {"output": str(result)}
            return {"output": str(client)}
        except Exception as e:
            return {"error": f"MCP tool call failed: {e}"}

    def list_resources(self, server_name: str | None = None) -> list[dict]:
        """List available MCP resources."""
        resources = []
        servers = [server_name] if server_name else list(self._connected)
        for name in servers:
            client = self._clients.get(name)
            if client and hasattr(client, "list_resources"):
                try:
                    result = client.list_resources()
                    for r in result.resources if hasattr(result, "resources") else []:
                        resources.append(
                            {
                                "server": name,
                                "uri": getattr(r, "uri", str(r)),
                                "name": getattr(r, "name", ""),
                                "description": getattr(r, "description", ""),
                            }
                        )
                except Exception as e:
                    logger.warning("Failed to list resources from %s: %s", name, e)
        return resources

    def get_status(self) -> dict[str, Any]:
        """Get status of all MCP servers."""
        status = {}
        for name, config in self._configs.items():
            status[name] = {
                "enabled": config.enabled,
                "connected": name in self._connected,
                "transport": config.transport,
                "tool_count": sum(
                    1 for t in self._tools.values() if t.server_name == name
                ),
            }
        return status

    # ── Internal ───────────────────────────────────────────────────────

    def _connect(self, name: str, config: MCPServerConfig) -> str | None:
        """Connect to a single MCP server. Returns error message or None."""
        if name in self._connected:
            return None

        try:
            # Try to import MCP SDK
            try:
                from mcp import ClientSession, StdioServerParameters  # noqa: F401
                from mcp.client.stdio import stdio_client  # noqa: F401
            except ImportError:
                logger.info("MCP SDK not installed. Skipping MCP server %s", name)
                return "mcp SDK not installed (pip install mcp)"

            if config.transport == "stdio":
                server_params = StdioServerParameters(
                    command=config.command or "",
                    args=config.args,
                    env=config.env or None,
                )

                # Run async connect in a new thread with event loop
                client_session, error = self._run_async_connect(name, server_params)
                if error or not client_session:
                    return error or f"Failed to connect to {name}"

                self._clients[name] = client_session
                self._connected.add(name)

                # Discover tools
                self._discover_tools(name, client_session)
                logger.info(
                    "MCP server '%s' connected with %d tools",
                    name,
                    sum(1 for t in self._tools.values() if t.server_name == name),
                )
                return None

            elif config.transport == "sse":
                return f"SSE transport not yet supported for {name}"

            return f"Unknown transport: {config.transport}"

        except Exception as e:
            logger.error("MCP connect error for %s: %s", name, e)
            return str(e)

    def _run_async_connect(self, name: str, server_params):
        """Run async MCP connection in a sync context."""
        from mcp import ClientSession
        from mcp.client.stdio import stdio_client

        result = {"session": None, "error": None}

        async def _connect_async():
            try:
                async with stdio_client(server_params) as (read, write):
                    async with ClientSession(read, write) as session:
                        await session.initialize()
                        result["session"] = _AsyncSessionWrapper(session)
                        # Keep the context managers alive by blocking
                        await asyncio.Future()  # Block forever
            except Exception as e:
                result["error"] = str(e)

        loop = asyncio.new_event_loop()
        thread = threading.Thread(
            target=loop.run_until_complete, args=(_connect_async(),), daemon=True
        )
        thread.start()
        thread.join(timeout=10)

        if result["error"] and not result["session"]:
            return None, result["error"]
        if result["session"]:
            return result["session"], None
        return None, "Connection timeout"

    def _discover_tools(self, server_name: str, client) -> None:
        """Discover and register tools from an MCP server."""
        # Remove old tools for this server
        old_prefix = f"mcp_{server_name}_"
        for k in list(self._tools.keys()):
            if k.startswith(old_prefix):
                del self._tools[k]

        # Discover new tools (sync wrapper)
        try:
            tools_result = client.list_tools() if hasattr(client, "list_tools") else None
            tools = tools_result.tools if hasattr(tools_result, "tools") else []

            for tool in tools:
                schema = {
                    "description": getattr(tool, "description", ""),
                    "inputSchema": getattr(
                        tool, "inputSchema", {"type": "object", "properties": {}}
                    ),
                    "annotations": getattr(tool, "annotations", {}),
                }
                tool_name = getattr(tool, "name", "unknown")

                bridge = MCPToolBridge(
                    server_name=server_name,
                    tool_name=tool_name,
                    tool_schema=schema,
                    call_fn=lambda sn=server_name, tn=tool_name, args=None:
                        self.call_tool(sn, tn, args or {}),
                )
                self._tools[bridge.name] = bridge
                logger.debug("Discovered MCP tool: %s", bridge.name)

        except Exception as e:
            logger.warning("Failed to discover tools from %s: %s", server_name, e)

    def _disconnect(self, name: str) -> None:
        """Disconnect from an MCP server."""
        self._clients.pop(name, None)
        self._connected.discard(name)
        # Remove tools
        prefix = f"mcp_{name}_"
        for k in list(self._tools.keys()):
            if k.startswith(prefix):
                del self._tools[k]


class _AsyncSessionWrapper:
    """Wraps an MCP ClientSession for sync access."""

    def __init__(self, session):
        self._session = session

    def __getattr__(self, name):
        attr = getattr(self._session, name)
        if asyncio.iscoroutinefunction(attr):

            def sync_wrapper(*args, **kwargs):
                try:
                    loop = asyncio.get_event_loop()
                    if loop.is_running():
                        # Can't await in running loop, return a placeholder
                        import concurrent.futures

                        with concurrent.futures.ThreadPoolExecutor() as pool:
                            future = pool.submit(asyncio.run, attr(*args, **kwargs))
                            return future.result(timeout=30)
                    else:
                        return loop.run_until_complete(attr(*args, **kwargs))
                except RuntimeError:
                    return asyncio.run(attr(*args, **kwargs))

            return sync_wrapper
        return attr


# ── Global Manager ─────────────────────────────────────────────────────

_global_manager: MCPClientManager | None = None


def get_mcp_manager() -> MCPClientManager:
    """Get or create the global MCP client manager."""
    global _global_manager
    if _global_manager is None:
        _global_manager = MCPClientManager()
    return _global_manager


__all__ = ["MCPToolBridge", "MCPClientManager", "get_mcp_manager"]