"""MCP client module for connecting to external MCP servers.

This module manages platform-wide MCP server connections that are
available to all agents as toolsets.

Example configuration in .env:
    MCP_SERVERS='[
        {"type": "stdio", "command": "uvx", "args": ["mcp-server-time"], "prefix": "time"},
        {"type": "sse", "url": "http://localhost:3001/sse", "prefix": "weather"},
        {"type": "http", "url": "http://localhost:8080/mcp", "prefix": "data"}
    ]'

Example usage:
    from pydanticai_multiagent.mcp.client import get_mcp_toolsets

    # Get all configured MCP servers as toolsets
    mcp_toolsets = await get_mcp_toolsets()

    # Use with an agent
    agent = Agent('openai:gpt-4o', toolsets=[*mcp_toolsets, my_other_toolset])
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from pydantic_ai.toolsets import AbstractToolset

logger = logging.getLogger(__name__)


@dataclass
class MCPServerConfig:
    """Configuration for a single MCP server."""

    type: str  # "stdio", "sse", or "http"
    prefix: str | None = None  # Tool prefix to avoid naming conflicts
    timeout: int = 30  # Timeout in seconds

    # For stdio servers
    command: str | None = None
    args: list[str] = field(default_factory=list)
    env: dict[str, str] | None = None

    # For SSE/HTTP servers
    url: str | None = None
    headers: dict[str, str] | None = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> MCPServerConfig:
        """Create config from dictionary."""
        return cls(
            type=data.get("type", "stdio"),
            prefix=data.get("prefix"),
            timeout=data.get("timeout", 30),
            command=data.get("command"),
            args=data.get("args", []),
            env=data.get("env"),
            url=data.get("url"),
            headers=data.get("headers"),
        )


class MCPClientManager:
    """Manages connections to external MCP servers.

    This class creates and manages MCP server connections based on
    configuration, providing them as toolsets for agents.
    """

    def __init__(self) -> None:
        self._servers: list[Any] = []
        self._initialized = False

    async def initialize(self, configs: list[dict[str, Any]]) -> None:
        """Initialize MCP server connections from configuration.

        Args:
            configs: List of MCP server configuration dictionaries.
        """
        if self._initialized:
            return

        for config_dict in configs:
            try:
                config = MCPServerConfig.from_dict(config_dict)
                server = self._create_server(config)
                if server:
                    self._servers.append(server)
                    logger.info(
                        f"Configured MCP server: {config.type} "
                        f"(prefix={config.prefix or 'none'})"
                    )
            except Exception as e:
                logger.error(f"Failed to configure MCP server: {e}")

        self._initialized = True

    def _create_server(self, config: MCPServerConfig) -> Any | None:
        """Create an MCP server instance from configuration.

        Args:
            config: Server configuration.

        Returns:
            MCP server instance or None if creation failed.
        """
        try:
            if config.type == "stdio":
                return self._create_stdio_server(config)
            elif config.type == "sse":
                return self._create_sse_server(config)
            elif config.type == "http":
                return self._create_http_server(config)
            else:
                logger.warning(f"Unknown MCP server type: {config.type}")
                return None
        except ImportError as e:
            logger.error(f"MCP dependencies not available: {e}")
            return None

    def _create_stdio_server(self, config: MCPServerConfig) -> Any:
        """Create a stdio-based MCP server."""
        from pydantic_ai.mcp import MCPServerStdio

        if not config.command:
            raise ValueError("stdio server requires 'command'")

        return MCPServerStdio(
            config.command,
            args=config.args,
            env=config.env,
            timeout=config.timeout,
            tool_prefix=config.prefix,
        )

    def _create_sse_server(self, config: MCPServerConfig) -> Any:
        """Create an SSE-based MCP server."""
        from pydantic_ai.mcp import MCPServerSSE

        if not config.url:
            raise ValueError("SSE server requires 'url'")

        return MCPServerSSE(
            config.url,
            headers=config.headers,
            timeout=config.timeout,
            tool_prefix=config.prefix,
        )

    def _create_http_server(self, config: MCPServerConfig) -> Any:
        """Create an HTTP-based MCP server (Streamable HTTP)."""
        from pydantic_ai.mcp import MCPServerStreamableHTTP

        if not config.url:
            raise ValueError("HTTP server requires 'url'")

        return MCPServerStreamableHTTP(
            config.url,
            headers=config.headers,
            timeout=config.timeout,
            tool_prefix=config.prefix,
        )

    def get_toolsets(self) -> list[AbstractToolset]:
        """Get all MCP servers as toolsets for agents.

        Returns:
            List of MCP server toolsets.
        """
        return list(self._servers)

    async def cleanup(self) -> None:
        """Clean up MCP server connections."""
        for server in self._servers:
            try:
                if hasattr(server, "close"):
                    await server.close()
                elif hasattr(server, "__aexit__"):
                    await server.__aexit__(None, None, None)
            except Exception as e:
                logger.error(f"Error closing MCP server: {e}")

        self._servers.clear()
        self._initialized = False


# Global manager instance
_manager = MCPClientManager()


async def initialize_mcp_clients(configs: list[dict[str, Any]] | None = None) -> None:
    """Initialize MCP client connections.

    Args:
        configs: Optional list of MCP server configs.
            If not provided, reads from settings.
    """
    if configs is None:
        from pydanticai_multiagent.config import settings
        configs = settings.mcp_servers

    if not configs:
        logger.debug("No MCP servers configured")
        return

    await _manager.initialize(configs)


def get_mcp_toolsets() -> list[AbstractToolset]:
    """Get all configured MCP servers as toolsets.

    Returns:
        List of MCP server toolsets to use with agents.

    Example:
        toolsets = get_mcp_toolsets()
        agent = Agent('openai:gpt-4o', toolsets=[*toolsets, my_toolset])
    """
    return _manager.get_toolsets()


async def cleanup_mcp_clients() -> None:
    """Clean up all MCP client connections."""
    await _manager.cleanup()
