"""MCP Server Registry — namespace-isolated lifecycle tracking.

Manages multiple :class:`~cleveragents.mcp.client.McpClient` instances
with independent lifecycles.  Each server is identified by a unique
namespace and can be started, stopped, and health-checked independently.

Usage::

    from cleveragents.mcp.registry import McpRegistry

    registry = McpRegistry()
    registry.register(client_config, namespace="github")
    registry.register(client_config_2, namespace="jira")

    # Clients start lazily on first tool call.
    result = registry.call_tool("github", "list_repos", {"org": "acme"})

    # Shut down all servers:
    registry.shutdown_all()
"""

from __future__ import annotations

import logging
import threading
from typing import Any

from cleveragents.mcp.adapter import MCPToolResult, MCPTransport
from cleveragents.mcp.client import McpClient, McpClientConfig

logger = logging.getLogger(__name__)

__all__ = [
    "McpRegistry",
]


class McpRegistry:
    """Registry of MCP clients indexed by namespace.

    Provides a unified interface for managing multiple MCP server
    connections with namespace isolation.
    """

    def __init__(self) -> None:
        self._clients: dict[str, McpClient] = {}
        self._transports: dict[str, MCPTransport | None] = {}
        self._lock = threading.RLock()

    def register(
        self,
        config: McpClientConfig,
        namespace: str,
        transport: MCPTransport | None = None,
    ) -> McpClient:
        """Register an MCP server configuration under a namespace.

        If a client already exists for the namespace, it is shut down
        and replaced.

        Parameters
        ----------
        config:
            Client configuration for the MCP server.
        namespace:
            Unique namespace identifier for the server.
        transport:
            Optional transport override (for testing).

        Returns
        -------
        McpClient
            The newly created client.
        """
        with self._lock:
            existing = self._clients.get(namespace)
            if existing is not None:
                existing.shutdown()

            client = McpClient(config=config, transport=transport)
            self._clients[namespace] = client
            self._transports[namespace] = transport

        logger.info(
            "Registered MCP server '%s' in namespace '%s'",
            config.server.name,
            namespace,
        )
        return client

    def unregister(self, namespace: str) -> None:
        """Unregister and shut down the client for a namespace.

        Parameters
        ----------
        namespace:
            Namespace to remove.

        Raises
        ------
        KeyError
            If no client is registered under the namespace.
        """
        with self._lock:
            client = self._clients.pop(namespace, None)
            self._transports.pop(namespace, None)

        if client is None:
            msg = f"No MCP client registered for namespace '{namespace}'"
            raise KeyError(msg)

        client.shutdown()
        logger.info("Unregistered MCP namespace '%s'", namespace)

    def get(self, namespace: str) -> McpClient | None:
        """Retrieve the client for a namespace, or ``None``."""
        with self._lock:
            return self._clients.get(namespace)

    def list_namespaces(self) -> list[str]:
        """Return all registered namespace names."""
        with self._lock:
            return list(self._clients.keys())

    def call_tool(
        self,
        namespace: str,
        tool_name: str,
        arguments: dict[str, Any],
    ) -> MCPToolResult:
        """Invoke a tool on a specific namespace's MCP server.

        The server is started lazily if not already running.

        Parameters
        ----------
        namespace:
            Namespace of the MCP server.
        tool_name:
            Name of the tool to invoke.
        arguments:
            Tool input arguments.

        Returns
        -------
        MCPToolResult
            Invocation outcome.

        Raises
        ------
        KeyError
            If no client is registered under the namespace.
        """
        with self._lock:
            client = self._clients.get(namespace)

        if client is None:
            msg = f"No MCP client registered for namespace '{namespace}'"
            raise KeyError(msg)

        return client.call_tool(tool_name, arguments)

    def shutdown_all(self) -> None:
        """Shut down all registered MCP clients."""
        with self._lock:
            clients = list(self._clients.values())

        for client in clients:
            try:
                client.shutdown()
            except Exception:
                logger.warning(
                    "Error shutting down MCP client '%s'",
                    client.server_name,
                    exc_info=True,
                )

        logger.info("All MCP clients shut down")

    def __len__(self) -> int:
        with self._lock:
            return len(self._clients)

    def __contains__(self, namespace: str) -> bool:
        with self._lock:
            return namespace in self._clients
