"""MCPProxy — manages the lifecycle of an external MCP server subprocess.

Wraps an MCP-compatible server process (e.g. ``uvx mcp-server-sqlite``)
started over *stdio* and exposes its tools through a synchronous Python API
that integrates naturally with the threaded callite server.

Example::

    proxy = MCPProxy("uvx", ["mcp-server-sqlite", "--db-path", "test.db"])
    proxy.connect()

    tools = proxy.get_tools()
    result = proxy.call_tool("list_tables", {})

    proxy.disconnect()

The class keeps a background ``asyncio`` event loop so that every public
method is safe to call from ordinary (non-async) threads — including the
daemon threads that ``RPCServer`` spawns for each incoming request.
"""

import asyncio
import logging
import threading
from contextlib import AsyncExitStack
from typing import Any

logger = logging.getLogger(__name__)


class MCPProxy:
    """Manages an external MCP server subprocess and exposes its tools.

    The proxy starts the subprocess lazily on the first ``connect()`` (or on
    the first ``get_tools`` / ``call_tool`` invocation) and keeps the
    connection alive for the lifetime of the object.  If the subprocess
    crashes, the next ``call_tool`` attempt will transparently reconnect.

    Args:
        command: Executable to start (e.g. ``"uvx"`` or ``"npx"``).
        args: Command-line arguments forwarded to the subprocess.
        env: Optional environment variables for the subprocess.
        call_timeout: Per-call timeout in seconds (``None`` = no timeout).
    """

    def __init__(
        self,
        command: str,
        args: list[str] | None = None,
        env: dict[str, str] | None = None,
        *,
        call_timeout: float | None = None,
    ):
        self.command = command
        self.args = args or []
        self.env = env
        self._call_timeout = call_timeout

        self._session: Any = None
        self._exit_stack: AsyncExitStack | None = None
        self._connected = False

        self._loop: asyncio.AbstractEventLoop | None = None
        self._loop_thread: threading.Thread | None = None
        self._connect_lock = threading.Lock()

        self._tools_cache: list[Any] | None = None

    # ------------------------------------------------------------------
    # Event-loop helpers
    # ------------------------------------------------------------------

    def _ensure_loop(self) -> None:
        """Spin up a background event loop if one is not already running."""
        with self._connect_lock:
            if self._loop is not None and self._loop.is_running():
                logger.debug("MCPProxy event loop already running")
                return
            logger.debug("MCPProxy starting event loop")
            self._loop = asyncio.new_event_loop()
            self._loop_thread = threading.Thread(
                target=self._loop.run_forever,
                daemon=True,
                name="mcp-proxy-loop",
            )
            self._loop_thread.start()

    def _run_async(self, coro: Any, timeout: float | None = None) -> Any:
        """Submit *coro* to the background loop and block until it finishes.

        Callers must ensure the event loop is running before calling this
        method (via ``_ensure_loop()``).  Do **not** call this while holding
        ``_connect_lock`` unless ``_ensure_loop`` was called *before* the lock
        was acquired — ``_ensure_loop`` itself acquires the same lock.
        """
        assert self._loop is not None, "Event loop not started; call _ensure_loop() first"
        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
        return future.result(timeout=timeout)

    # ------------------------------------------------------------------
    # Async internals
    # ------------------------------------------------------------------

    async def _connect_async(self) -> None:
        try:
            from mcp import ClientSession, StdioServerParameters
            from mcp.client.stdio import stdio_client
        except ImportError:
            raise ImportError(
                "MCPProxy requires the 'mcp' package. "
                "Install it with: pip install ..."
            )

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

        self._exit_stack = AsyncExitStack()
        read, write = await self._exit_stack.enter_async_context(
            stdio_client(server_params)
        )
        self._session = await self._exit_stack.enter_async_context(
            ClientSession(read, write)
        )
        await self._session.initialize()
        self._connected = True
        logger.info(
            "MCPProxy connected to %s %s", self.command, " ".join(self.args)
        )

    async def _disconnect_async(self) -> None:
        if self._exit_stack:
            try:
                await self._exit_stack.aclose()
            except Exception:
                pass
        self._session = None
        self._exit_stack = None
        self._connected = False
        self._tools_cache = None

    async def _reconnect_async(self) -> None:
        await self._disconnect_async()
        await self._connect_async()

    async def _get_tools_async(self) -> list[Any]:
        assert self._session is not None
        result = await self._session.list_tools()
        return result.tools

    async def _list_tools_async(self) -> list[Any]:
        return await self._get_tools_async()

    async def _list_prompts_async(self) -> list[Any]:
        assert self._session is not None
        result = await self._session.list_prompts()
        return result.prompts

    async def _get_prompt_async(self, name: str, arguments: dict[str, str] | None = None) -> Any:
        assert self._session is not None
        return await self._session.get_prompt(name, arguments)

    async def _list_resources_async(self) -> list[Any]:
        assert self._session is not None
        result = await self._session.list_resources()
        return result.resources

    async def _read_resource_async(self, uri: str) -> Any:
        assert self._session is not None
        from pydantic import AnyUrl
        return await self._session.read_resource(AnyUrl(uri))

    async def _list_resource_templates_async(self) -> list[Any]:
        assert self._session is not None
        result = await self._session.list_resource_templates()
        return result.resourceTemplates

    async def _send_ping_async(self) -> Any:
        assert self._session is not None
        return await self._session.send_ping()

    async def _call_tool_async(
        self, name: str, arguments: dict[str, Any]
    ) -> Any:
        assert self._session is not None
        result = await self._session.call_tool(name, arguments=arguments)

        if result.isError:
            error_text = "\n".join(
                c.text for c in result.content if hasattr(c, "text")
            )
            raise RuntimeError(f"MCP tool '{name}' error: {error_text}")

        # Extract text content from the result
        texts = [c.text for c in result.content if hasattr(c, "text")]
        if len(texts) == 1:
            return texts[0]
        return texts if texts else None

    # ------------------------------------------------------------------
    # Public synchronous API
    # ------------------------------------------------------------------

    def connect(self) -> None:
        """Start the subprocess and perform the MCP handshake.

        Safe to call multiple times — subsequent calls are no-ops while the
        connection is alive.
        """
        logger.debug("MCPProxy connecting to %s", self.command)
        self._ensure_loop()
        with self._connect_lock:
            if self._connected:
                return
            self._run_async(self._connect_async())

    def _ensure_connected(self) -> None:
        if not self._connected:
            self.connect()

    def get_tools(self) -> list[Any]:
        """Return the list of tools published by the external MCP server.

        Results are cached after the first call.  Use ``refresh_tools()`` to
        force a re-fetch.
        """
        logger.debug("MCPProxy getting tools from %s", self.command)
        self._ensure_connected()
        if self._tools_cache is None:
            self._tools_cache = self._run_async(self._get_tools_async())
        return self._tools_cache

    def refresh_tools(self) -> list[Any]:
        """Re-fetch the tool list from the subprocess (clears the cache)."""
        self._tools_cache = None
        return self.get_tools()

    def call_tool(
        self, name: str, arguments: dict[str, Any] | None = None
    ) -> Any:
        """Execute a tool on the external MCP server.

        If the subprocess connection is lost, one automatic reconnect attempt
        is made before the error is propagated.

        Args:
            name: Tool name as published by the MCP server.
            arguments: Keyword arguments expected by the tool.

        Returns:
            The text content returned by the tool.  If there are multiple
            text blocks the return value is a ``list[str]``.
        """
        self._ensure_connected()
        try:
            return self._run_async(
                self._call_tool_async(name, arguments or {}),
                timeout=self._call_timeout,
            )
        except (ConnectionError, BrokenPipeError, OSError, EOFError):
            logger.warning(
                "MCPProxy connection to %s lost, reconnecting…",
                self.command,
            )
            with self._connect_lock:
                self._run_async(self._reconnect_async())
            return self._run_async(
                self._call_tool_async(name, arguments or {}),
                timeout=self._call_timeout,
            )

    def list_tools(self) -> list[Any]:
        """Return the list of tools via the MCP protocol ``list_tools`` method."""
        self._ensure_connected()
        return self._run_async(self._list_tools_async())

    def list_prompts(self) -> list[Any]:
        """Return the list of prompts via the MCP protocol ``list_prompts`` method."""
        self._ensure_connected()
        return self._run_async(self._list_prompts_async())

    def get_prompt(self, name: str, arguments: dict[str, str] | None = None) -> Any:
        """Retrieve a prompt by name via the MCP protocol ``get_prompt`` method."""
        self._ensure_connected()
        return self._run_async(self._get_prompt_async(name, arguments))

    def list_resources(self) -> list[Any]:
        """Return the list of resources via the MCP protocol ``list_resources`` method."""
        self._ensure_connected()
        return self._run_async(self._list_resources_async())

    def read_resource(self, uri: str) -> Any:
        """Read a resource by URI via the MCP protocol ``read_resource`` method."""
        self._ensure_connected()
        return self._run_async(self._read_resource_async(uri))

    def list_resource_templates(self) -> list[Any]:
        """Return the list of resource templates via the MCP protocol."""
        self._ensure_connected()
        return self._run_async(self._list_resource_templates_async())

    def send_ping(self) -> Any:
        """Send a ping via the MCP protocol ``send_ping`` method."""
        self._ensure_connected()
        return self._run_async(self._send_ping_async())

    def disconnect(self) -> None:
        """Stop the subprocess and release all resources."""
        with self._connect_lock:
            if self._loop and self._connected:
                try:
                    self._run_async(self._disconnect_async())
                except Exception:
                    pass
            self._session = None
            self._exit_stack = None
            self._connected = False
            self._tools_cache = None
            if self._loop:
                self._loop.call_soon_threadsafe(self._loop.stop)
                self._loop = None
                self._loop_thread = None
            logger.info("MCPProxy disconnected from %s", self.command)
