# 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

"""Manager Facade for the MCP Gateway.

This module provides the `McpManagerAsync` class (the core) and `McpManager` class (the sync facade),
which serve as the primary entry points and coordinators for the CoReason MCP Gateway.
"""

from typing import Any, Dict, List, Optional, Union

import anyio
import httpx
from mcp.types import CallToolResult, Prompt, Resource

from coreason_mcp.aggregator import ToolAggregator
from coreason_mcp.config import McpConfig, McpServerConfig
from coreason_mcp.invocation_proxy import InvocationProxy
from coreason_mcp.registry import ServerRegistry
from coreason_mcp.session_manager import SessionManager
from coreason_mcp.types import ServerUrlMapping, ToolMapping
from coreason_mcp.utils.logger import logger


class McpManagerAsync:
    """Facade for managing MCP connections, tools, and invocations (Async).

    Acts as the high-level interface to the MCP Gateway, coordinating the
    registry, session manager, aggregator, and invocation proxy.

    Attributes:
        config: The global configuration for the gateway.
        session_manager: Handles creation of client sessions.
        registry: Manages active connections and retries.
        aggregator: Aggregates tools and resources from sessions.
        invocation_proxy: Routes tool calls to the correct session.
    """

    def __init__(
        self,
        config: Union[McpConfig, List[McpServerConfig]],
        client: Optional[httpx.AsyncClient] = None,
    ) -> None:
        """Initialize the McpManagerAsync.

        Args:
            config: Either an McpConfig object or a list of McpServerConfig objects.
            client: Optional external httpx.AsyncClient for connection pooling.
        """
        self._internal_client = client is None
        self.client = client or httpx.AsyncClient()

        if isinstance(config, list):
            self.config = McpConfig(servers=config)
        else:
            self.config = config

        server_configs = {server.name: server for server in self.config.servers}
        self.session_manager = SessionManager(server_configs=server_configs)
        self.registry = ServerRegistry(self.session_manager, timeout=self.config.timeout)
        self.aggregator = ToolAggregator()
        self.invocation_proxy = InvocationProxy()

        self._tool_mapping: ToolMapping = {}
        self._server_urls: ServerUrlMapping = {}

        # Populate server URLs
        for server in self.config.servers:
            target = server.url if server.url else f"Command: {server.command}"
            self._server_urls[server.name] = target

    async def __aenter__(self) -> "McpManagerAsync":
        await self.connect_all()
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.close()
        if self._internal_client:
            await self.client.aclose()

    async def connect_all(self) -> None:
        """Connects to all configured MCP servers.

        Initiates the "Connect" phase, establishing persistent async connections
        via SSE or Stdio transports.
        """
        logger.info("Initializing connection to all MCP servers...")
        await self.registry.connect_all(self.config.servers)

    async def get_openai_tools(self) -> List[Dict[str, Any]]:
        """Aggregates tools from all connected servers and converts them to OpenAI format.

        Performs the "Aggregate" and "Translate" phases, merging tools into a unified namespace.

        Returns:
            A list of tool definitions in OpenAI JSON schema format.
        """
        active_sessions = self.registry.get_active_sessions()
        llm_tools, mapping = await self.aggregator.aggregate_tools(active_sessions)
        self._tool_mapping = mapping
        return llm_tools

    async def list_resources(self) -> List[Resource]:
        """Aggregates resources from all connected servers.

        Returns:
            A list of Resource objects.
        """
        active_sessions = self.registry.get_active_sessions()
        return await self.aggregator.aggregate_resources(active_sessions)

    async def list_prompts(self) -> List[Prompt]:
        """Aggregates prompts from all connected servers.

        Returns:
            A list of Prompt objects.
        """
        active_sessions = self.registry.get_active_sessions()
        return await self.aggregator.aggregate_prompts(active_sessions)

    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> CallToolResult:
        """Executes a tool by name.

        Performs the "Route" phase, delegating to the InvocationProxy to find the
        correct server and execute the tool.

        Args:
            tool_name: The namespaced name of the tool to execute.
            arguments: The arguments for the tool.

        Returns:
            The result of the tool execution as a CallToolResult object.
        """
        active_sessions = self.registry.get_active_sessions()
        return await self.invocation_proxy.call_tool(
            tool_name, arguments, active_sessions, self._tool_mapping, self._server_urls
        )

    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
        clean termination of all server connections.
        """
        logger.info("Closing all MCP connections...")
        await self.registry.close()
        self._tool_mapping.clear()


class McpManager:
    """Synchronous Facade for McpManagerAsync.

    Provides a blocking interface to the MCP Gateway, wrapping the async core
    using anyio.run.
    """

    def __init__(self, config: Union[McpConfig, List[McpServerConfig]]) -> None:
        """Initialize the McpManager facade.

        Args:
            config: Either an McpConfig object or a list of McpServerConfig objects.
        """
        self._async = McpManagerAsync(config)

    def __enter__(self) -> "McpManager":
        return self

    def __exit__(self, *args: Any) -> None:
        async def _wrapper() -> None:
            await self._async.__aexit__(*args)

        anyio.run(_wrapper)

    def connect_all(self) -> None:
        """Connects to all configured MCP servers."""
        anyio.run(self._async.connect_all)

    def get_openai_tools(self) -> List[Dict[str, Any]]:
        """Aggregates tools from all connected servers."""
        return anyio.run(self._async.get_openai_tools)

    def list_resources(self) -> List[Resource]:
        """Aggregates resources from all connected servers."""
        return anyio.run(self._async.list_resources)

    def list_prompts(self) -> List[Prompt]:
        """Aggregates prompts from all connected servers."""
        return anyio.run(self._async.list_prompts)

    def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> CallToolResult:
        """Executes a tool by name."""
        return anyio.run(self._async.call_tool, tool_name, arguments)

    def close(self) -> None:
        """Closes all active connections."""
        anyio.run(self._async.close)
