# SPDX-License-Identifier: Apache-2.0
"""
MCP Client Manager for handling multiple MCP server connections.
"""

import asyncio
import logging
from typing import Any, Dict, List, Optional

from .client import MCPClient
from .tools import merge_tools, mcp_tools_to_openai, openai_call_to_mcp
from .types import (
    MCPConfig,
    MCPServerStatus,
    MCPTool,
    MCPToolResult,
)

logger = logging.getLogger(__name__)


class MCPClientManager:
    """
    Manages multiple MCP server connections.

    Provides a unified interface for:
    - Connecting to multiple MCP servers
    - Discovering and aggregating tools
    - Executing tool calls
    - Managing connection lifecycle
    """

    def __init__(self, config: MCPConfig):
        """
        Initialize MCP Client Manager.

        Args:
            config: MCP configuration with server definitions
        """
        self.config = config
        self._clients: Dict[str, MCPClient] = {}
        self._started = False
        self._lock = asyncio.Lock()

        # Create clients for each server
        for name, server_config in config.servers.items():
            self._clients[name] = MCPClient(server_config)

    @property
    def is_started(self) -> bool:
        """Check if manager has been started."""
        return self._started

    async def start(self):
        """
        Start the manager and connect to all enabled servers.

        Connections are made in parallel for faster startup.
        """
        async with self._lock:
            if self._started:
                return

            logger.info(
                f"Starting MCP client manager with {len(self._clients)} servers"
            )

            # Connect to all servers in parallel
            tasks = [
                client.connect()
                for client in self._clients.values()
                if client.config.enabled
            ]

            if tasks:
                results = await asyncio.gather(*tasks, return_exceptions=True)

                # Log results
                for client, result in zip(
                    [c for c in self._clients.values() if c.config.enabled],
                    results,
                ):
                    if isinstance(result, Exception):
                        logger.error(f"Failed to connect to '{client.name}': {result}")
                    elif result:
                        logger.info(f"Connected to '{client.name}'")

            self._started = True

            # Log summary
            connected = sum(1 for c in self._clients.values() if c.is_connected)
            total_tools = sum(len(c.tools) for c in self._clients.values())
            logger.info(
                f"MCP manager started: {connected}/{len(self._clients)} servers, "
                f"{total_tools} tools available"
            )

    async def stop(self):
        """Stop the manager and disconnect from all servers."""
        async with self._lock:
            if not self._started:
                return

            logger.info("Stopping MCP client manager")

            # Disconnect from all servers in parallel
            tasks = [client.disconnect() for client in self._clients.values()]
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)

            self._started = False
            logger.info("MCP client manager stopped")

    def get_all_tools(self) -> List[MCPTool]:
        """
        Get all tools from all connected servers.

        Returns:
            List of MCPTool instances
        """
        tools = []
        for client in self._clients.values():
            if client.is_connected:
                tools.extend(client.tools)
        return tools

    def get_all_tools_openai(self) -> List[Dict[str, Any]]:
        """
        Get all tools in OpenAI function calling format.

        Returns:
            List of OpenAI-compatible tool definitions
        """
        return mcp_tools_to_openai(self.get_all_tools())

    def get_merged_tools(
        self,
        user_tools: Optional[List[Dict[str, Any]]] = None,
    ) -> List[Dict[str, Any]]:
        """
        Get MCP tools merged with user-provided tools.

        User tools take precedence on name conflicts.

        Args:
            user_tools: Optional user-provided tools in OpenAI format

        Returns:
            Combined list of tools in OpenAI format
        """
        return merge_tools(self.get_all_tools(), user_tools)

    def get_server_status(self) -> List[MCPServerStatus]:
        """
        Get status of all servers.

        Returns:
            List of MCPServerStatus for each server
        """
        return [client.get_status() for client in self._clients.values()]

    def get_client(self, server_name: str) -> Optional[MCPClient]:
        """
        Get client for a specific server.

        Args:
            server_name: Name of the server

        Returns:
            MCPClient instance or None if not found
        """
        return self._clients.get(server_name)

    async def execute_tool(
        self,
        full_name: str,
        arguments: Dict[str, Any],
        timeout: Optional[float] = None,
    ) -> MCPToolResult:
        """
        Execute a tool by its full name (server__tool).

        Args:
            full_name: Full tool name with server prefix
            arguments: Tool arguments
            timeout: Optional timeout in seconds

        Returns:
            MCPToolResult with the result or error
        """
        # Parse full name
        server_name, tool_name, _ = openai_call_to_mcp(
            {"function": {"name": full_name, "arguments": "{}"}}
        )

        # If no server prefix, try to find the tool
        if not server_name:
            server_name = self._find_tool_server(full_name)
            tool_name = full_name

        if not server_name:
            return MCPToolResult(
                tool_name=full_name,
                content=None,
                is_error=True,
                error_message=f"Tool '{full_name}' not found in any connected server",
            )

        # Get client
        client = self._clients.get(server_name)
        if not client:
            return MCPToolResult(
                tool_name=full_name,
                content=None,
                is_error=True,
                error_message=f"Server '{server_name}' not found",
            )

        if not client.is_connected:
            return MCPToolResult(
                tool_name=full_name,
                content=None,
                is_error=True,
                error_message=f"Server '{server_name}' is not connected",
            )

        # Execute tool
        return await client.call_tool(
            tool_name,
            arguments,
            timeout=timeout or self.config.default_timeout,
        )

    async def execute_tool_call(
        self,
        tool_call: Dict[str, Any],
        timeout: Optional[float] = None,
    ) -> MCPToolResult:
        """
        Execute a tool call from OpenAI format.

        Args:
            tool_call: OpenAI tool call object
            timeout: Optional timeout in seconds

        Returns:
            MCPToolResult with the result or error
        """
        server_name, tool_name, arguments = openai_call_to_mcp(tool_call)

        if server_name:
            full_name = f"{server_name}__{tool_name}"
        else:
            full_name = tool_name

        return await self.execute_tool(full_name, arguments, timeout)

    def _find_tool_server(self, tool_name: str) -> Optional[str]:
        """
        Find which server has a tool by name.

        Args:
            tool_name: Tool name (without server prefix)

        Returns:
            Server name or None if not found
        """
        for client in self._clients.values():
            if client.is_connected:
                for tool in client.tools:
                    if tool.name == tool_name:
                        return client.name
        return None

    async def refresh_tools(self):
        """Refresh tools from all connected servers."""
        tasks = [
            client.refresh_tools()
            for client in self._clients.values()
            if client.is_connected
        ]
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    async def reconnect(self, server_name: Optional[str] = None):
        """
        Reconnect to server(s).

        Args:
            server_name: Specific server to reconnect, or None for all
        """
        if server_name:
            client = self._clients.get(server_name)
            if client:
                await client.disconnect()
                await client.connect()
        else:
            # Reconnect all
            for client in self._clients.values():
                await client.disconnect()
                await client.connect()
