# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_mcp

"""Registry for managing active MCP server connections.

This module provides the `ServerRegistry` class, which handles the lifecycle
of connections to multiple MCP servers, including connection establishment,
retries, and cleanup.
"""

import asyncio
from contextlib import AsyncExitStack
from typing import Dict, List

from mcp import ClientSession

from coreason_mcp.config import McpServerConfig
from coreason_mcp.session_manager import SessionManager
from coreason_mcp.utils.logger import logger


class ServerRegistry:
    """Manages a registry of active MCP server connections.

    Responsible for maintaining the "Connect" phase of the loop. Handles connection
    lifecycle, including initial connection, retries with backoff, and health checks.

    Attributes:
        session_manager: The SessionManager instance used for creating connections.
        max_retries: Maximum number of connection attempts per server.
        retry_delay: Delay in seconds between retries.
        timeout: Default timeout for requests in seconds.
    """

    def __init__(
        self,
        session_manager: SessionManager,
        max_retries: int = 3,
        retry_delay: float = 2.0,
        timeout: float = 30.0,
    ) -> None:
        """Initialize the ServerRegistry.

        Args:
            session_manager: The SessionManager instance to use for creating connections.
            max_retries: Maximum number of connection attempts per server. Defaults to 3.
            retry_delay: Delay in seconds between retries. Defaults to 2.0.
            timeout: Default timeout for requests in seconds. Defaults to 30.0.
        """
        self.session_manager = session_manager
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.timeout = timeout
        self._exit_stack = AsyncExitStack()
        self._active_sessions: Dict[str, ClientSession] = {}
        self._lock = asyncio.Lock()

    async def connect_all(self, server_configs: List[McpServerConfig]) -> None:
        """Connects to all provided server configurations with retry logic.

        Iterates through the provided configurations and attempts to establish
        a connection to each server. Connections are handled sequentially to
        ensure proper context management.

        Args:
            server_configs: List of server configurations to connect to.
        """
        logger.info("ServerRegistry: Initializing connections...")
        # We must connect sequentially to ensure the context manager is entered
        # in the same task that will eventually close it (the main task).
        # Using asyncio.TaskGroup here would cause 'anyio' RuntimeErrors during shutdown.
        for config in server_configs:
            await self._connect_with_retry(config)

    async def _connect_with_retry(self, config: McpServerConfig) -> None:
        """Attempts to connect to a single server with retries.

        Args:
            config: The server configuration.
        """
        attempt = 0
        while attempt <= self.max_retries:
            try:
                attempt += 1
                target = config.url if config.url else f"Command: {config.command}"
                logger.info(
                    f"Connecting to MCP Server: {config.name} at {target} (Attempt {attempt}/{self.max_retries + 1})"
                )

                # Enter context manager and get session
                # Lock is required because AsyncExitStack.enter_async_context is not thread/task safe
                # when modifying the internal stack structure concurrently.
                async with self._lock:
                    session = await self._exit_stack.enter_async_context(
                        self.session_manager.connect(config, timeout=self.timeout)
                    )

                # Perform Health Check
                logger.debug(f"Pinging {config.name}...")
                await session.send_ping()

                self._active_sessions[config.name] = session
                logger.info(f"Connected to {config.name}")
                return  # Success

            except Exception as e:
                logger.warning(f"Failed to connect to {config.name} (Attempt {attempt}): {e}")
                if attempt <= self.max_retries:
                    logger.info(f"Retrying in {self.retry_delay} seconds...")
                    await asyncio.sleep(self.retry_delay)
                else:
                    logger.error(f"Max retries reached for {config.name}. Skipping.")

    def get_active_sessions(self) -> Dict[str, ClientSession]:
        """Returns a dictionary of active sessions.

        Returns:
            A dictionary mapping server names to connected ClientSession objects.
        """
        return self._active_sessions

    async def close(self) -> None:
        """Closes all active connections and cleans up resources.

        This method should be called when shutting down the gateway to ensure
        all server connections are properly terminated.
        """
        logger.info("ServerRegistry: Closing all connections...")
        async with self._lock:
            await self._exit_stack.aclose()
        self._active_sessions.clear()
