"""MCP connection — lifecycle manager for a single MCP server.

Wraps the ``mcp`` SDK's ``ClientSession`` + transport setup into a
simple async context manager that handles connect / list / call / close.

Supports two transports:

- **stdio** — spawn a subprocess, communicate over stdin/stdout
- **url** — connect to a running server via SSE / Streamable HTTP
"""

from __future__ import annotations

import logging
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

try:
    from mcp.client.stdio import stdio_client

    from mcp import ClientSession, StdioServerParameters

    _HAS_MCP = True
except ImportError:  # optional dep not installed
    _HAS_MCP = False

if TYPE_CHECKING:
    from mcp.types import CallToolResult, Tool

logger = logging.getLogger(__name__)


def _require_mcp() -> None:
    """Raise a clear error when the ``mcp`` package is missing."""
    if not _HAS_MCP:
        raise ImportError(
            "The 'mcp' package is required for MCP client support. "
            "Install it with: pip install albusos[mcp]"
        )


@dataclass
class McpConnection:
    """Manages the lifecycle of one MCP server connection.

    Parameters
    ----------
    namespace
        Prefix for registered tool names (e.g. ``"servicem8"``).
    command
        Command + args to spawn an stdio-based server
        (e.g. ``["python", "-m", "servicem8_mcp"]``).
    url
        URL of a running MCP server (SSE / Streamable HTTP).
    env
        Optional environment variables passed to the subprocess.

    Usage
    -----
    Typically created via :func:`infrastructure.mcp.adapter.connect_mcp`
    rather than directly.
    """

    namespace: str
    command: list[str] | None = None
    url: str | None = None
    env: dict[str, str] | None = None

    # Internal state — not part of the public API
    _session: ClientSession | None = field(default=None, repr=False, init=False)
    _exit_stack: AsyncExitStack = field(default_factory=AsyncExitStack, repr=False, init=False)
    _tools: list[Tool] = field(default_factory=list, repr=False, init=False)
    _connected: bool = field(default=False, repr=False, init=False)

    async def connect(self) -> list[Tool]:
        """Spawn the server, initialise the session, and discover tools.

        Returns
        -------
        list[Tool]
            Tools advertised by the MCP server.

        Raises
        ------
        ImportError
            If the ``mcp`` package is not installed.
        ValueError
            If neither *command* nor *url* is provided.
        RuntimeError
            If the connection fails or the server does not respond.
        """
        _require_mcp()

        if self._connected:
            return self._tools

        if self.command:
            await self._connect_stdio()
        elif self.url:
            await self._connect_url()
        else:
            raise ValueError("Either 'command' (stdio) or 'url' (HTTP) must be provided.")

        # Initialise the MCP session handshake
        assert self._session is not None
        await self._session.initialize()

        # Discover tools
        response = await self._session.list_tools()
        self._tools = list(response.tools)
        self._connected = True

        logger.info(
            "MCP [%s] connected — %d tools: %s",
            self.namespace,
            len(self._tools),
            [t.name for t in self._tools],
        )
        return self._tools

    async def call_tool(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
        """Invoke a tool on the connected MCP server.

        Parameters
        ----------
        name
            Tool name *as known to the MCP server* (without namespace prefix).
        arguments
            Tool arguments dict.

        Returns
        -------
        CallToolResult
            Raw MCP result.
        """
        if not self._connected or self._session is None:
            raise RuntimeError(f"MCP [{self.namespace}] not connected — call connect() first.")
        return await self._session.call_tool(name, arguments)

    def list_tools(self) -> list[Tool]:
        """Return the cached list of tools discovered at connect time."""
        return list(self._tools)

    async def close(self) -> None:
        """Tear down the transport and session."""
        if self._connected:
            await self._exit_stack.aclose()
            self._session = None
            self._tools = []
            self._connected = False
            logger.info("MCP [%s] disconnected.", self.namespace)

    @property
    def connected(self) -> bool:
        """Whether the connection is active."""
        return self._connected

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    async def _connect_stdio(self) -> None:
        """Set up an stdio transport."""
        assert self.command
        server_params = StdioServerParameters(
            command=self.command[0],
            args=self.command[1:],
            env=self.env,
        )
        transport = await self._exit_stack.enter_async_context(stdio_client(server_params))
        read_stream, write_stream = transport
        self._session = await self._exit_stack.enter_async_context(
            ClientSession(read_stream, write_stream),
        )

    async def _connect_url(self) -> None:
        """Set up an SSE / Streamable HTTP transport."""
        assert self.url
        # The MCP SDK provides sse_client for SSE and streamablehttp_client
        # for the newer Streamable HTTP transport.  Try streamablehttp first
        # (the modern default), fall back to SSE for older servers.
        try:
            from mcp.client.streamable_http import streamablehttp_client

            transport = await self._exit_stack.enter_async_context(
                streamablehttp_client(self.url),
            )
        except (ImportError, Exception):
            from mcp.client.sse import sse_client

            transport = await self._exit_stack.enter_async_context(sse_client(self.url))

        read_stream, write_stream = transport[0], transport[1]
        self._session = await self._exit_stack.enter_async_context(
            ClientSession(read_stream, write_stream),
        )


__all__ = ["McpConnection"]
