"""MCP (Model Context Protocol) client — connects to MCP servers and converts
their tools to yier_agents Tool instances."""

from __future__ import annotations

import contextlib
import sys
from pathlib import Path
from collections.abc import AsyncIterator
from typing import Any

from pydantic import BaseModel, ConfigDict

from yier_agents.src.config import McpServerDict
from yier_agents.src.log import logger
from yier_agents.src.tool import Tool, ToolContext, ToolOutput

_SKIP_STATUSES = {"disabled", "failed", "needs_auth", "needs_client_registration"}


def _make_params_model(tool_name: str, input_schema: dict) -> type[BaseModel]:
    """Create a Pydantic model that preserves an MCP tool's JSON schema.

    The model accepts arbitrary extra fields (to pass through any arguments)
    and returns the original MCP inputSchema from model_json_schema() so that
    Tool.to_openai_schema() produces the correct schema for the LLM.
    """
    schema = dict(input_schema)

    class MCPToolParams(BaseModel):
        model_config = ConfigDict(extra="allow")

        @classmethod
        def model_json_schema(cls, **kwargs: Any) -> dict:  # type: ignore[override]
            return schema

    MCPToolParams.__name__ = f"{tool_name}_params"
    MCPToolParams.__qualname__ = f"{tool_name}_params"
    return MCPToolParams


class MCPClient:
    """Async context manager that connects to an MCP server and exposes its
    tools as yier_agents Tool instances.

    Usage (stdio)::

        async with MCPClient.stdio("uvx", ["mcp-server-fetch"]) as client:
            tools = await client.get_tools()
            agent = Agent(llm=llm, tools=tools)

    Usage (SSE — legacy GET-based stream)::

        async with MCPClient.sse("http://localhost:8000/sse") as client:
            tools = await client.get_tools()
            agent = Agent(llm=llm, tools=tools)

    Usage (Streamable HTTP — modern POST-based, use when SSE returns 405)::

        async with MCPClient.http("http://localhost:8000/mcp") as client:
            tools = await client.get_tools()
            agent = Agent(llm=llm, tools=tools)
    """

    def __init__(self, transport: str, **kwargs: Any) -> None:
        self._transport = transport
        self._kwargs = kwargs
        self._session: Any = None
        self._exit_stack = contextlib.AsyncExitStack()

    # ── factory methods ───────────────────────────────────────────────────────

    @classmethod
    def stdio(
        cls,
        command: str,
        args: list[str],
        env: dict[str, str] | None = None,
    ) -> MCPClient:
        """Connect via stdio transport (spawns a subprocess)."""
        return cls("stdio", command=command, args=args, env=env)

    @classmethod
    def sse(
        cls,
        url: str,
        headers: dict[str, str] | None = None,
    ) -> MCPClient:
        """Connect via SSE transport (legacy GET-based HTTP stream).

        Use this for servers that expose a /sse endpoint.
        If you get 405 Method Not Allowed, try MCPClient.http() instead.
        """
        return cls("sse", url=url, headers=headers)

    @classmethod
    def http(
        cls,
        url: str,
        headers: dict[str, str] | None = None,
    ) -> MCPClient:
        """Connect via Streamable HTTP transport (modern POST-based).

        Use this for servers that return 405 when using MCPClient.sse().
        This is the newer MCP transport (spec 2025-03-26).
        """
        return cls("http", url=url, headers=headers)

    @classmethod
    def from_server_config(cls, server_cfg: McpServerDict) -> MCPClient:
        """Build one client from a typed config entry."""
        transport_type = server_cfg.get("type")
        if transport_type == "stdio":
            return cls.stdio(
                command=server_cfg["command"],
                args=server_cfg.get("args", []),
                env=server_cfg.get("env"),
            )

        if transport_type == "http":
            return cls.http(
                url=server_cfg["url"],
                headers=server_cfg.get("headers"),
            )

        if transport_type == "sse":
            return cls.sse(
                url=server_cfg["url"],
                headers=server_cfg.get("headers"),
            )

        raise ValueError(f"Unknown MCP server type: {transport_type}")

    @classmethod
    @contextlib.asynccontextmanager
    async def from_config(
        cls,
        config_dir: Path | None = None,
    ) -> AsyncIterator[dict[str, MCPClient]]:
        """Async context manager that loads and connects all active MCP servers
        from the yier config file.

        Skips servers whose status is disabled, failed, needs_auth, or
        needs_client_registration. Servers with no status field are included.

        Usage::

            async with MCPClient.from_config() as clients:
                tools = []
                for client in clients.values():
                    tools.extend(await client.get_tools())
        """
        from yier_agents.src.config import YIERConfig

        mcp_configs = YIERConfig.read_mcp_config(config_dir)
        unconnected: dict[str, MCPClient] = {}

        for name, server_cfg in mcp_configs.items():
            status = server_cfg.get("status")
            if status in _SKIP_STATUSES:
                logger.debug(f"Skipping MCP server '{name}' (status={status})")
                continue

            transport_type = server_cfg.get("type")
            try:
                unconnected[name] = cls.from_server_config(server_cfg)
            except ValueError:
                logger.warning(f"Unknown MCP server type '{transport_type}' for '{name}', skipping.")

        async with contextlib.AsyncExitStack() as stack:
            connected = {
                name: await stack.enter_async_context(client)
                for name, client in unconnected.items()
            }
            logger.debug(f"Connected {len(connected)} MCP client(s): {list(connected)}")
            yield connected

    # ── context manager ───────────────────────────────────────────────────────

    async def __aenter__(self) -> MCPClient:
        from mcp import ClientSession, StdioServerParameters
        from mcp.client.stdio import stdio_client
        from mcp.client.sse import sse_client

        await self._exit_stack.__aenter__()
        try:
            if self._transport == "stdio":
                server_params = StdioServerParameters(
                    command=self._kwargs["command"],
                    args=self._kwargs.get("args", []),
                    env=self._kwargs.get("env"),
                )
                read, write = await self._exit_stack.enter_async_context(
                    stdio_client(server_params)
                )
            elif self._transport == "http":
                from mcp.client.streamable_http import streamable_http_client
                import httpx

                http_client = httpx.AsyncClient(
                    headers=self._kwargs.get("headers") or {}
                )
                read, write, _ = await self._exit_stack.enter_async_context(
                    streamable_http_client(
                        url=self._kwargs["url"],
                        http_client=http_client,
                    )
                )
            else:  # sse
                read, write = await self._exit_stack.enter_async_context(
                    sse_client(
                        url=self._kwargs["url"],
                        headers=self._kwargs.get("headers"),
                    )
                )

            self._session = await self._exit_stack.enter_async_context(
                ClientSession(read, write)
            )
            await self._session.initialize()
            return self
        except BaseException:
            await self._exit_stack.__aexit__(*sys.exc_info())
            raise

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
        self._session = None

    # ── public API ────────────────────────────────────────────────────────────

    async def get_tools(self) -> list[Tool]:
        """Return all tools on the MCP server as Tool instances."""
        if self._session is None:
            raise RuntimeError("MCPClient must be used as an async context manager")
        result = await self._session.list_tools()
        return [self._convert_tool(t) for t in result.tools]

    # ── internal ──────────────────────────────────────────────────────────────

    def _convert_tool(self, mcp_tool: Any) -> Tool:
        """Convert one MCP tool definition to a Tool instance."""
        params_model = _make_params_model(
            mcp_tool.name,
            mcp_tool.inputSchema or {},
        )
        session = self._session

        async def execute(params: BaseModel, ctx: ToolContext) -> ToolOutput:
            arguments = params.model_dump(exclude_unset=True)
            call_result = await session.call_tool(mcp_tool.name, arguments)

            text_parts = [c.text for c in call_result.content if hasattr(c, "text")]
            content_str = "\n".join(text_parts)

            if call_result.isError:
                raise RuntimeError(content_str)

            return ToolOutput(content=content_str)

        return Tool(
            name=mcp_tool.name,
            description=mcp_tool.description or "",
            parameters=params_model,
            execute=execute,
        )
