"""MCP Client Service - Connect to external MCP servers.

Uses the official `mcp` package for protocol handling.
Supports stdio transport for local subprocess servers.

Example:
    client = McpClientService([
        McpServerConfig(
            id="github",
            command="npx",
            args=["-y", "@modelcontextprotocol/server-github"],
            env={"GITHUB_TOKEN": "..."},
        ),
    ])

    await client.connect_all()
    tools = await client.list_tools("github")
    result = await client.call_tool("github", "search_repos", {"query": "python"})
    await client.close_all()
"""

from __future__ import annotations

import logging
from typing import Any

from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


class McpServerConfig(BaseModel):
    """Configuration for an MCP server connection."""

    id: str
    command: str
    args: list[str] = Field(default_factory=list)
    env: dict[str, str] = Field(default_factory=dict)
    timeout_s: float = 30.0


class McpTool(BaseModel):
    """Tool exposed by an MCP server."""

    name: str
    description: str = ""
    input_schema: dict[str, Any] = Field(default_factory=dict)


class McpClientService:
    """Manages connections to multiple MCP servers.

    Uses the official `mcp` package for protocol handling.
    """

    def __init__(self, servers: list[McpServerConfig] | None = None):
        self._configs: dict[str, McpServerConfig] = {}
        self._sessions: dict[str, Any] = {}  # server_id -> ClientSession
        self._contexts: dict[str, Any] = {}  # server_id -> context manager

        for config in servers or []:
            self._configs[config.id] = config

    def add_server(self, config: McpServerConfig) -> None:
        """Add a server configuration."""
        self._configs[config.id] = config

    def list_server_ids(self) -> list[str]:
        """List configured server IDs."""
        return list(self._configs.keys())

    async def connect(self, server_id: str) -> bool:
        """Connect to a specific MCP server."""
        if server_id in self._sessions:
            return True  # Already connected

        config = self._configs.get(server_id)
        if not config:
            logger.warning(f"No config for MCP server: {server_id}")
            return False

        try:
            from mcp import ClientSession
            from mcp.client.stdio import stdio_client

            # Build command
            command = config.command
            args = config.args
            env = config.env or None

            # Create stdio connection
            server_params = {
                "command": command,
                "args": args,
            }
            if env:
                server_params["env"] = env

            # Connect using context manager pattern
            read_stream, write_stream = await stdio_client(**server_params).__aenter__()
            session = ClientSession(read_stream, write_stream)
            await session.__aenter__()

            # Initialize the session
            await session.initialize()

            self._sessions[server_id] = session
            logger.info(f"Connected to MCP server: {server_id}")
            return True

        except ImportError:
            logger.exception("MCP package not installed. Run: pip install mcp")
            return False
        except Exception as e:
            logger.exception(f"Failed to connect to MCP server {server_id}: {e}")
            return False

    async def connect_all(self) -> dict[str, bool]:
        """Connect to all configured servers."""
        results = {}
        for server_id in self._configs:
            results[server_id] = await self.connect(server_id)
        return results

    async def list_tools(self, server_id: str) -> list[McpTool]:
        """List tools available from a server."""
        session = self._sessions.get(server_id)
        if not session:
            # Try to connect
            if not await self.connect(server_id):
                return []
            session = self._sessions.get(server_id)

        if not session:
            return []

        try:
            result = await session.list_tools()
            tools = []
            for tool in result.tools:
                tools.append(
                    McpTool(
                        name=tool.name,
                        description=tool.description or "",
                        input_schema=tool.inputSchema or {},
                    )
                )
            return tools
        except Exception as e:
            logger.exception(f"Failed to list tools from {server_id}: {e}")
            return []

    async def call_tool(
        self,
        server_id: str,
        tool: str,
        arguments: dict[str, Any] | None = None,
    ) -> Any:
        """Call a tool on an MCP server."""
        session = self._sessions.get(server_id)
        if not session:
            if not await self.connect(server_id):
                msg = f"Cannot connect to MCP server: {server_id}"
                raise RuntimeError(msg)
            session = self._sessions.get(server_id)

        if not session:
            msg = f"No session for MCP server: {server_id}"
            raise RuntimeError(msg)

        try:
            result = await session.call_tool(tool, arguments or {})

            # Extract content from result
            if hasattr(result, "content") and result.content:
                # MCP returns content as a list of content blocks
                contents = []
                for block in result.content:
                    if hasattr(block, "text"):
                        contents.append(block.text)
                    elif hasattr(block, "data"):
                        contents.append(block.data)
                    else:
                        contents.append(str(block))
                return contents[0] if len(contents) == 1 else contents

            return result

        except Exception as e:
            logger.exception(f"MCP tool call failed {server_id}.{tool}: {e}")
            raise

    async def close(self, server_id: str) -> None:
        """Close connection to a specific server."""
        session = self._sessions.pop(server_id, None)
        if session:
            try:
                await session.__aexit__(None, None, None)
            except Exception as e:
                logger.warning(f"Error closing MCP session {server_id}: {e}")

    async def close_all(self) -> None:
        """Close all connections."""
        for server_id in list(self._sessions.keys()):
            await self.close(server_id)

    def is_connected(self, server_id: str) -> bool:
        """Check if connected to a server."""
        return server_id in self._sessions


__all__ = [
    "McpClientService",
    "McpServerConfig",
    "McpTool",
]
