"""MCP client — connects to MCP servers and exposes their tools."""
from __future__ import annotations

from typing import Any

from logos_adk.circuit_breaker import CircuitBreaker, CircuitBreakerConfig
from logos_adk.tools import Tool


class MCPClient:
    """Client that connects to an MCP server and imports its tools.

    Supports two transport modes:
    - **stdio**: Launch a subprocess (e.g., ``command="npx @modelcontextprotocol/server-github"``)
    - **sse**: Connect to an HTTP SSE endpoint (e.g., ``url="http://localhost:3000/sse"``)

    Usage::

        # stdio — launch a CLI tool as MCP server
        client = MCPClient(command="npx @modelcontextprotocol/server-github")
        await client.connect()
        tools = client.get_tools()  # list[Tool]

        # SSE — connect to a running server
        client = MCPClient(url="http://localhost:3000/sse")
        await client.connect()
        tools = client.get_tools()

        # Shorthand from string
        client = MCPClient.from_string("npx @modelcontextprotocol/server-github")
        client = MCPClient.from_string("http://localhost:3000/sse")
    """

    def __init__(
        self,
        *,
        command: str | None = None,
        url: str | None = None,
        env: dict[str, str] | None = None,
        headers: dict[str, str] | None = None,
        timeout: float = 30.0,
    ) -> None:
        if not command and not url:
            raise ValueError("MCPClient requires either 'command' (stdio) or 'url' (SSE)")
        if command and url:
            raise ValueError("MCPClient accepts either 'command' or 'url', not both")

        self._command = command
        self._url = url
        self._env = env
        self._headers = headers or {}
        self._timeout = timeout
        self._tools: list[Tool] = []
        self._session: Any = None
        self._transport: Any = None
        self._connected = False
        self._circuit_breaker = CircuitBreaker(
            name=f"mcp:{command or url}",
            config=CircuitBreakerConfig(
                failure_threshold=5,
                reset_timeout_seconds=30.0,
            ),
        )

    @classmethod
    def from_string(cls, source: str, **kwargs: Any) -> MCPClient:
        """Create an MCPClient from a string — auto-detect stdio vs SSE."""
        if source.startswith("http://") or source.startswith("https://"):
            return cls(url=source, **kwargs)
        return cls(command=source, **kwargs)

    async def connect(self) -> None:
        """Connect to the MCP server and discover tools."""
        try:
            from mcp import ClientSession
            from mcp.client.stdio import stdio_client, StdioServerParameters
            from mcp.client.sse import sse_client
        except ImportError:
            raise ImportError(
                "mcp package required for MCP support. Install with: "
                "pip install mcp"
            )

        if self._command:
            await self._connect_stdio(ClientSession, StdioServerParameters, stdio_client)
        else:
            await self._connect_sse(ClientSession, sse_client)

    async def _connect_stdio(self, ClientSession: Any, StdioServerParameters: Any, stdio_client: Any) -> None:
        """Connect via stdio transport."""
        import shlex
        import os

        parts = shlex.split(self._command)
        env = {**os.environ, **(self._env or {})}

        server_params = StdioServerParameters(
            command=parts[0],
            args=parts[1:] if len(parts) > 1 else [],
            env=env,
        )

        self._transport = stdio_client(server_params)
        read_stream, write_stream = await self._transport.__aenter__()
        self._session = ClientSession(read_stream, write_stream)
        await self._session.__aenter__()
        await self._session.initialize()
        self._connected = True
        await self._discover_tools()

    async def _connect_sse(self, ClientSession: Any, sse_client: Any) -> None:
        """Connect via SSE transport."""
        self._transport = sse_client(
            url=self._url,
            headers=self._headers,
            timeout=self._timeout,
        )
        read_stream, write_stream = await self._transport.__aenter__()
        self._session = ClientSession(read_stream, write_stream)
        await self._session.__aenter__()
        await self._session.initialize()
        self._connected = True
        await self._discover_tools()

    async def _discover_tools(self) -> None:
        """List tools from the MCP server and convert to Logos Tool objects."""
        result = await self._session.list_tools()
        self._tools = []
        for mcp_tool in result.tools:
            tool = self._wrap_mcp_tool(mcp_tool)
            self._tools.append(tool)

    def _wrap_mcp_tool(self, mcp_tool: Any) -> Tool:
        """Wrap an MCP tool definition as a Logos Tool."""
        session = self._session
        tool_name = mcp_tool.name
        description = mcp_tool.description or ""
        input_schema = mcp_tool.inputSchema or {"type": "object", "properties": {}}
        breaker = self._circuit_breaker

        async def _call_mcp(**kwargs: Any) -> str:
            async with breaker:
                result = await session.call_tool(tool_name, arguments=kwargs)
            # MCP returns content as a list of content blocks
            parts = []
            for block in result.content:
                if hasattr(block, "text"):
                    parts.append(block.text)
                else:
                    parts.append(str(block))
            return "\n".join(parts)

        tool = Tool(fn=_call_mcp, name=tool_name, description=description)
        # Override auto-generated schema with the MCP-provided schema
        tool._schema = {
            "name": tool_name,
            "description": description,
            "parameters": input_schema,
        }
        return tool

    def get_tools(self) -> list[Tool]:
        """Return discovered tools. Must call connect() first."""
        if not self._connected:
            raise RuntimeError("MCPClient not connected. Call await client.connect() first.")
        return list(self._tools)

    async def disconnect(self) -> None:
        """Disconnect from the MCP server."""
        if self._session:
            try:
                await self._session.__aexit__(None, None, None)
            except Exception:
                pass
        if self._transport:
            try:
                await self._transport.__aexit__(None, None, None)
            except Exception:
                pass
        self._connected = False
        self._tools = []

    async def __aenter__(self) -> MCPClient:
        await self.connect()
        return self

    async def __aexit__(self, *exc: Any) -> None:
        await self.disconnect()

    @property
    def connected(self) -> bool:
        return self._connected

    def __repr__(self) -> str:
        source = self._command or self._url
        status = "connected" if self._connected else "disconnected"
        return f"MCPClient({source!r}, {status}, tools={len(self._tools)})"
