"""MCP client manager — persistent connections to MCP servers.

Reads ~/.icarus/mcp.json (same format as Claude Code's mcp.json),
connects to each server at agent startup, and holds connections
open for the session lifetime.  Supports both stdio and HTTP
(streamable-http) transports.

Isolation guarantees:
- Each server gets its own AsyncExitStack — one crash can't affect others.
- Tools are wrapped in proxies that catch all errors (including
  BaseExceptionGroup from anyio) so nothing leaks into the agent.
"""

from __future__ import annotations

import asyncio
import functools
import logging
from contextlib import AsyncExitStack
from pathlib import Path
from typing import Any, Callable

logger = logging.getLogger(__name__)

OnServerFailure = Callable[[str, str], Any]    # (server_name, error_msg) -> None
OnServerRecovery = Callable[[str], Any]        # (server_name) -> None
OnServerReconnect = Callable[[str], Any]       # (server_name) -> None


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

async def _check_reachable(url: str, timeout: float = 5) -> None:
    """Fast TCP check that an HTTP MCP server is reachable.

    Raises ConnectionError immediately if the port is not listening.
    This prevents entering streamablehttp_client which hangs forever
    on dead connections (session.initialize() blocks and ignores
    asyncio cancellation because it uses anyio internally).
    """
    import socket
    from urllib.parse import urlparse

    parsed = urlparse(url)
    host = parsed.hostname or "localhost"
    port = parsed.port or (443 if parsed.scheme == "https" else 80)

    def _probe() -> None:
        sock = socket.create_connection((host, port), timeout=timeout)
        sock.close()

    try:
        await asyncio.get_running_loop().run_in_executor(None, _probe)
    except (OSError, socket.timeout) as exc:
        raise ConnectionError(f"Cannot reach {host}:{port}") from exc


async def _open_session(cfg: dict[str, Any], stack: AsyncExitStack) -> Any:
    """Connect to a single MCP server and return an initialized ClientSession.

    Dispatches to stdio or streamable-http based on the presence of a
    ``url`` key in *cfg*.  The transport and session are entered into
    *stack* so the caller controls their lifetime.
    """
    from mcp import ClientSession

    if "url" in cfg:
        await _check_reachable(cfg["url"])
        from mcp.client.streamable_http import streamablehttp_client

        read, write, _ = await stack.enter_async_context(
            streamablehttp_client(url=cfg["url"], headers=cfg.get("headers"))
        )
    else:
        from mcp import StdioServerParameters
        from mcp.client.stdio import stdio_client

        params = StdioServerParameters(
            command=cfg["command"],
            args=cfg.get("args", []),
            env=cfg.get("env"),
        )
        read, write = await stack.enter_async_context(stdio_client(params))

    session = await stack.enter_async_context(ClientSession(read, write))
    await session.initialize()
    return session


async def _safe_aclose(stack: AsyncExitStack) -> None:
    """Close an exit stack, suppressing anyio task-group teardown errors."""
    try:
        await stack.aclose()
    except (Exception, BaseExceptionGroup):
        logger.warning("MCP: error during cleanup", exc_info=True)


def _extract_leaf_error(exc: BaseException) -> BaseException:
    """Walk nested ExceptionGroups and return the first non-group exception."""
    while isinstance(exc, BaseExceptionGroup) and exc.exceptions:
        exc = exc.exceptions[0]
    return exc


def _wrap_tool(
    tool: Any,
    server_name: str,
    server_cfg: dict[str, Any],
    server_state: dict[str, bool],
    on_failure: OnServerFailure | None,
    on_recovery: OnServerRecovery | None,
    on_reconnect: OnServerReconnect | None,
) -> Any:
    """Wrap a LangChain MCP tool so errors are caught and returned as text.

    *server_state* is shared across all tools from the same server so
    callbacks fire exactly once per state transition.
    """
    original_arun = tool._arun

    def _unavailable(error_msg: str):
        msg = f"[Tool unavailable] {tool.name} — connector '{server_name}' is offline: {error_msg}"
        if getattr(tool, "response_format", None) == "content_and_artifact":
            return msg, None
        return msg

    @functools.wraps(original_arun)
    async def safe_arun(*args: Any, **kwargs: Any) -> str:
        # If reconnect is already in progress, short-circuit.
        if server_state.get("reconnecting"):
            return _unavailable("Reconnecting — try again in a moment")

        # If degraded and we have a URL, do a quick TCP check.
        # If server is back, trigger reconnect and short-circuit.
        # If server is still down, short-circuit (don't hang on dead tool).
        if server_state.get("degraded") and on_reconnect:
            url = server_cfg.get("url")
            if url:
                # Set flag BEFORE the await to prevent concurrent tools
                # from all firing on_reconnect simultaneously.
                server_state["reconnecting"] = True
                try:
                    await _check_reachable(url, timeout=2)
                    try:
                        on_reconnect(server_name)
                    except Exception:
                        pass
                    return _unavailable("Reconnecting — try again in a moment")
                except ConnectionError:
                    server_state["reconnecting"] = False
                    return _unavailable("Server still offline")

        try:
            task = asyncio.create_task(original_arun(*args, **kwargs))
            done, pending = await asyncio.wait({task}, timeout=30)
            if pending:
                task.cancel()
                raise TimeoutError("MCP tool call timed out after 30s")
            result = task.result()
            if server_state.get("degraded") and on_recovery:
                server_state["degraded"] = False
                try:
                    on_recovery(server_name)
                except Exception:
                    pass
            return result
        except (Exception, BaseExceptionGroup, asyncio.CancelledError) as exc:
            leaf = _extract_leaf_error(exc) if isinstance(exc, BaseExceptionGroup) else exc
            error_msg = str(leaf) or type(leaf).__name__
            logger.warning("MCP tool '%s' (server '%s') failed: %s", tool.name, server_name, error_msg)
            if on_failure and not server_state.get("degraded"):
                server_state["degraded"] = True
                try:
                    on_failure(server_name, error_msg)
                except Exception:
                    pass
            return _unavailable(error_msg)

    tool._arun = safe_arun
    tool._server_name = server_name
    return tool


# ---------------------------------------------------------------------------
# MCPManager
# ---------------------------------------------------------------------------

class MCPManager:
    """Manages persistent MCP server connections for an agent session.

    Each server gets its own AsyncExitStack for isolation — a crash in one
    server's connection can never affect another server or the sidecar.
    Tools are wrapped in proxies that catch all errors so nothing leaks
    into the agent graph.
    """

    def __init__(
        self,
        servers: dict[str, dict[str, Any]],
        disabled: dict[str, dict[str, Any]] | None = None,
        on_failure: OnServerFailure | None = None,
        on_recovery: OnServerRecovery | None = None,
        on_reconnect: OnServerReconnect | None = None,
    ) -> None:
        self._servers = servers
        self._disabled = disabled or {}
        self._on_failure = on_failure
        self._on_recovery = on_recovery
        self._on_reconnect = on_reconnect
        # Per-server isolation
        self._stacks: dict[str, AsyncExitStack] = {}
        self._tools: list[Any] = []
        # Per-server tracking for the UI
        self._server_status: dict[str, str] = {}
        self._server_tools: dict[str, list[str]] = {}
        self._server_errors: dict[str, str] = {}

        for name in self._disabled:
            self._server_status[name] = "disabled"
            self._server_tools[name] = []

    async def start(self) -> tuple[list[Any], list[str]]:
        """Connect to all configured MCP servers and load their tools.

        Returns:
            (tools, failed) — combined tool list and names of servers
            that failed to connect.
        """
        from langchain_mcp_adapters.tools import load_mcp_tools

        failed: list[str] = []

        for name, cfg in self._servers.items():
            stack = AsyncExitStack()
            try:
                # Enter context from a separate task so anyio's cancel
                # scope is on THAT task, not the caller's.  When the
                # server dies, the scope targets the finished task.
                session = await asyncio.create_task(_open_session(cfg, stack))
                tools = await load_mcp_tools(session)
                state: dict[str, bool] = {}
                wrapped = [_wrap_tool(t, name, cfg, state, self._on_failure, self._on_recovery, self._on_reconnect) for t in tools]
                self._tools.extend(wrapped)
                self._stacks[name] = stack
                self._server_status[name] = "connected"
                self._server_tools[name] = [t.name for t in tools]
                logger.info("MCP: connected to '%s' (%d tools)", name, len(tools))
            except (Exception, BaseExceptionGroup) as exc:
                leaf = _extract_leaf_error(exc) if isinstance(exc, BaseExceptionGroup) else exc
                self._server_status[name] = "error"
                self._server_errors[name] = str(leaf)
                self._server_tools[name] = []
                failed.append(name)
                logger.warning("MCP: failed to connect to '%s', skipping", name, exc_info=True)
                # Fire-and-forget: _safe_aclose hangs if streamablehttp_client
                # context was entered but the connection failed mid-way.
                asyncio.ensure_future(_safe_aclose(stack))

        connected = self.server_names
        if connected:
            logger.info(
                "MCP: %d tools loaded from %d server(s): %s",
                len(self._tools), len(connected), ", ".join(connected),
            )
        return self._tools, failed

    async def stop(self) -> None:
        """Disconnect all MCP servers independently."""
        for name, stack in list(self._stacks.items()):
            await _safe_aclose(stack)
        self._stacks.clear()
        self._tools.clear()

    def mark_degraded(self, name: str, error: str) -> None:
        """Mark a server as degraded — connected but tools are failing."""
        if self._server_status.get(name) == "connected":
            self._server_status[name] = "degraded"
            self._server_errors[name] = error

    def mark_recovered(self, name: str) -> None:
        """Flip a degraded server back to connected."""
        if self._server_status.get(name) == "degraded":
            self._server_status[name] = "connected"
            self._server_errors.pop(name, None)

    async def disconnect_server(self, name: str) -> None:
        """Tear down a single server's connection."""
        stack = self._stacks.pop(name, None)
        if stack:
            await _safe_aclose(stack)
        self._tools = [t for t in self._tools if not getattr(t, '_server_name', None) == name]
        self._server_status[name] = "disabled"

    @staticmethod
    async def probe(cfg: dict[str, Any]) -> list[str]:
        """Try connecting to a single server config. Return tool names or raise."""
        from langchain_mcp_adapters.tools import load_mcp_tools

        stack = AsyncExitStack()
        conn_error: BaseException | None = None
        try:
            # Wrap in create_task to isolate anyio's cancel scope —
            # same pattern as start() (line 240).  Use asyncio.wait
            # for timeout because session.initialize() ignores
            # asyncio cancellation (anyio internals).
            task = asyncio.create_task(_open_session(cfg, stack))
            done, pending = await asyncio.wait({task}, timeout=30)
            if pending:
                task.cancel()
                raise TimeoutError("MCP probe timed out after 30s")
            session = task.result()
            tools = await load_mcp_tools(session)
            return [t.name for t in tools]
        except BaseException as exc:
            conn_error = exc
            raise
        finally:
            try:
                await stack.aclose()
            except (Exception, BaseExceptionGroup) as cleanup_exc:
                logger.warning("MCP: error closing probe connection", exc_info=True)
                if isinstance(conn_error, asyncio.CancelledError):
                    leaf = _extract_leaf_error(cleanup_exc)
                    raise ConnectionError(str(leaf)) from leaf

    # -- Properties ----------------------------------------------------------

    @property
    def tool_count(self) -> int:
        return len(self._tools)

    @property
    def server_names(self) -> list[str]:
        return [n for n, s in self._server_status.items() if s == "connected"]

    @property
    def server_status(self) -> dict[str, str]:
        """Per-server status: ``"connected"``, ``"error"``, or ``"disabled"``."""
        return dict(self._server_status)

    @property
    def server_tools(self) -> dict[str, list[str]]:
        return dict(self._server_tools)

    @property
    def server_errors(self) -> dict[str, str]:
        return dict(self._server_errors)

    # -- Factory -------------------------------------------------------------

    @classmethod
    def from_config_file(
        cls,
        config_path: Path,
        on_failure: OnServerFailure | None = None,
        on_recovery: OnServerRecovery | None = None,
        on_reconnect: OnServerReconnect | None = None,
    ) -> MCPManager | None:
        """Load MCP config from a JSON file.

        Returns MCPManager instance, or None if no servers configured.
        """
        from icarus.mcp.config import read_mcp_config

        config = read_mcp_config(config_path)
        servers = config.get("mcpServers", {})
        disabled = config.get("_disabled", {})

        if not servers and not disabled:
            return None

        logger.info("MCP: found %d enabled + %d disabled server(s) in %s",
                     len(servers), len(disabled), config_path)
        return cls(servers, disabled, on_failure=on_failure, on_recovery=on_recovery, on_reconnect=on_reconnect)
