"""Persistent MCP manager with config-driven reload support."""

from __future__ import annotations

import asyncio
import contextlib
import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Literal

from yier_agents.src.config import McpServerDict, YIERConfig
from yier_agents.src.log import logger
from yier_agents.src.mcp_client import MCPClient
from yier_agents.src.tool import Tool

RuntimeStatus = Literal[
    "connected",
    "disabled",
    "failed",
    "needs_auth",
    "needs_client_registration",
]


@dataclass(slots=True)
class MCPServerRuntimeState:
    """Runtime state for one configured MCP server."""

    name: str
    config: McpServerDict | None = None
    fingerprint: str = ""
    status: RuntimeStatus = "disabled"
    error: str | None = None
    client: MCPClient | None = None
    tools: list[Tool] = field(default_factory=list)

    def snapshot(self) -> dict[str, Any]:
        """Return a serializable status snapshot."""
        payload: dict[str, Any] = {
            "status": self.status,
            "tool_count": len(self.tools),
        }
        if self.error:
            payload["error"] = self.error
        return payload


class MCPManager:
    """Keep MCP connections alive and reload them when config changes."""

    def __init__(
        self,
        config_dir: Path | None = None,
        poll_interval: float = 1.0,
        client_factory: Callable[[McpServerDict], MCPClient] | None = None,
    ) -> None:
        self.config_dir = config_dir
        self.poll_interval = poll_interval
        self.config_path = YIERConfig.resolve_config_path(config_dir)
        self.client_factory = client_factory or MCPClient.from_server_config
        self._states: dict[str, MCPServerRuntimeState] = {}
        self._lock = asyncio.Lock()
        self._watch_task: asyncio.Task[None] | None = None
        self._version = 0
        self._config_marker = self._read_config_marker()

    @property
    def version(self) -> int:
        """Return a monotonically increasing version for state changes."""
        return self._version

    async def __aenter__(self) -> MCPManager:
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.stop()

    async def start(self) -> None:
        """Load config, connect active servers, and start watching for changes."""
        await self.reload(force_reconnect=True)
        self.start_watching()

    async def stop(self) -> None:
        """Stop the watcher and disconnect every managed client."""
        watch_task = self._watch_task
        self._watch_task = None
        if watch_task is not None:
            watch_task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await watch_task

        async with self._lock:
            states = list(self._states.values())
            self._states.clear()
            self._version += 1

        for state in states:
            await self._close_client(state.client)

    def start_watching(self) -> None:
        """Start the background config watcher if it is not already running."""
        if self._watch_task is not None and not self._watch_task.done():
            return
        self._watch_task = asyncio.create_task(self.watch_config(), name="mcp-config-watch")

    async def watch_config(self) -> None:
        """Poll the config file and reload when it changes."""
        while True:
            await asyncio.sleep(self.poll_interval)
            await self.reload_if_changed()

    async def reload_if_changed(self) -> bool:
        """Reload active MCP servers if the config file changed."""
        current_marker = self._read_config_marker()
        if current_marker == self._config_marker:
            return False

        self._config_marker = current_marker
        await self.reload()
        return True

    async def reload(self, force_reconnect: bool = False) -> None:
        """Apply config changes and reconnect only the affected servers."""
        desired = YIERConfig.read_mcp_config(self.config_dir)

        async with self._lock:
            existing_names = set(self._states)
            desired_names = set(desired)

            for removed_name in existing_names - desired_names:
                state = self._states.pop(removed_name)
                await self._close_client(state.client)
                self._version += 1

            for name, server_cfg in desired.items():
                enabled = server_cfg.get("enabled", True)
                fingerprint = json.dumps(server_cfg, sort_keys=True, ensure_ascii=False)
                current = self._states.get(name)

                if not enabled or server_cfg.get("status") == "disabled":
                    if current is not None:
                        await self._close_client(current.client)
                    self._states[name] = MCPServerRuntimeState(
                        name=name,
                        config=server_cfg,
                        fingerprint=fingerprint,
                        status="disabled",
                    )
                    self._version += 1
                    continue

                if (
                    not force_reconnect
                    and current is not None
                    and current.fingerprint == fingerprint
                    and current.client is not None
                    and current.status == "connected"
                ):
                    continue

                if current is not None:
                    await self._close_client(current.client)

                self._states[name] = await self._connect_server(
                    name=name,
                    server_cfg=server_cfg,
                    fingerprint=fingerprint,
                )
                self._version += 1

    async def refresh_tools(self) -> bool:
        """Refresh tools from already-connected servers."""
        changed = False
        async with self._lock:
            for state in self._states.values():
                if state.client is None or state.status != "connected":
                    continue

                try:
                    tools = await state.client.get_tools()
                except Exception as exc:
                    logger.warning(f"Failed to refresh MCP tools for '{state.name}': {exc}")
                    await self._close_client(state.client)
                    state.client = None
                    state.tools = []
                    state.status = "failed"
                    state.error = str(exc)
                    changed = True
                    continue

                if self._tool_signature(state.tools) != self._tool_signature(tools):
                    state.tools = tools
                    state.error = None
                    changed = True

            if changed:
                self._version += 1

        return changed

    async def get_tools(self, refresh: bool = False) -> list[Tool]:
        """Return tools from every connected server."""
        if refresh:
            await self.refresh_tools()

        async with self._lock:
            result: list[Tool] = []
            for state in self._states.values():
                if state.status == "connected":
                    result.extend(state.tools)
            return result

    async def get_clients(self) -> dict[str, MCPClient]:
        """Return currently connected clients keyed by server name."""
        async with self._lock:
            return {
                name: state.client
                for name, state in self._states.items()
                if state.client is not None and state.status == "connected"
            }

    async def get_status(self) -> dict[str, dict[str, Any]]:
        """Return a serializable snapshot of current MCP runtime state."""
        async with self._lock:
            return {
                name: state.snapshot()
                for name, state in sorted(self._states.items())
            }

    async def _connect_server(
        self,
        name: str,
        server_cfg: McpServerDict,
        fingerprint: str,
    ) -> MCPServerRuntimeState:
        client = self.client_factory(server_cfg)
        try:
            await client.__aenter__()
            tools = await client.get_tools()
        except Exception as exc:
            await self._close_client(client)
            logger.warning(f"Failed to connect MCP server '{name}': {exc}")
            return MCPServerRuntimeState(
                name=name,
                config=server_cfg,
                fingerprint=fingerprint,
                status="failed",
                error=str(exc),
            )

        return MCPServerRuntimeState(
            name=name,
            config=server_cfg,
            fingerprint=fingerprint,
            status="connected",
            client=client,
            tools=tools,
        )

    async def _close_client(self, client: MCPClient | None) -> None:
        """Close one client and ignore cleanup errors."""
        if client is None:
            return

        try:
            await client.__aexit__(None, None, None)
        except Exception as exc:
            logger.debug(f"Failed to close MCP client cleanly: {exc}")

    def _read_config_marker(self) -> tuple[bool, int, int]:
        """Build a cheap marker for detecting config file changes."""
        try:
            stat = self.config_path.stat()
        except FileNotFoundError:
            return (False, 0, 0)
        return (True, stat.st_mtime_ns, stat.st_size)

    @staticmethod
    def _tool_signature(tools: list[Tool]) -> tuple[tuple[str, str], ...]:
        """Build a stable signature for the current tool set."""
        return tuple(
            sorted((tool.name, tool.description) for tool in tools)
        )
