"""MCP Client setup and management for LangGraph ReAct Agent."""

from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
from typing import Any, Callable

from langchain_mcp_adapters.client import (  # type: ignore[import-untyped]
    MultiServerMCPClient,
)

logger = logging.getLogger(__name__)


@dataclass
class MCPValidationResult:
    """Result of MCP server validation."""

    success: bool
    tools: list[dict[str, Any]] = field(default_factory=list)
    error: str | None = None
    server_details: dict[str, Any] = field(default_factory=dict)


class MCPAdapter:
    """Adapter for managing MCP client, tools and validation."""

    # Supported transport types
    SUPPORTED_TRANSPORTS = ("stdio", "streamable_http", "sse")

    def __init__(self) -> None:
        """Initialize the MCP adapter."""
        self._client: MultiServerMCPClient | None = None
        self._tools_cache: list[Callable[..., object]] = []
        self._server_configs: str | None = None

    def validate(self, mcp_server_configs: dict[str, Any]) -> MCPValidationResult:
        """Validate MCP server configurations.

        Args:
            mcp_server_configs: Dictionary of server name to server config.

        Returns:
            MCPValidationResult with validation status and details.
        """
        if not mcp_server_configs:
            return MCPValidationResult(
                success=False,
                error="No MCP server configurations provided",
            )

        errors: list[str] = []
        valid_servers: list[str] = []

        for server_name, config in mcp_server_configs.items():
            if not isinstance(config, dict):
                errors.append(f"Server '{server_name}': config must be a dictionary")
                continue

            transport = config.get("transport", "stdio")

            if transport == "stdio":
                if "command" not in config:
                    errors.append(
                        f"Server '{server_name}': missing 'command' field for stdio transport"
                    )
                    continue
            elif transport in ("streamable_http", "sse"):
                if "url" not in config:
                    errors.append(
                        f"Server '{server_name}': missing 'url' field for {transport} transport"
                    )
                    continue
            else:
                errors.append(
                    f"Server '{server_name}': unsupported transport '{transport}'"
                )
                continue

            valid_servers.append(server_name)

        if errors:
            return MCPValidationResult(
                success=False,
                error="; ".join(errors),
                server_details={
                    "server_count": len(mcp_server_configs),
                    "valid_servers": valid_servers,
                    "invalid_count": len(errors),
                },
            )

        return MCPValidationResult(
            success=True,
            tools=[],
            server_details={
                "server_count": len(mcp_server_configs),
                "tool_count": 0,
                "servers": valid_servers,
            },
        )

    async def get_client(
        self,
        server_configs: str,
    ) -> MultiServerMCPClient | None:
        """Get or initialize the MCP client with given server configurations.

        Args:
            server_configs: JSON string of server configurations.

        Returns:
            MultiServerMCPClient instance or None if initialization fails.
        """
        if self._client is None or self._server_configs != server_configs:
            try:
                config = json.loads(server_configs)
            except Exception as e:
                logger.error("Invalid JSON for server_configs: %s", e)
                return None

            # Validate configuration before initializing client
            validation_result = self.validate(config)
            if not validation_result.success:
                logger.error(
                    "MCP config validation failed: %s", validation_result.error
                )
                return None

            try:
                self._client = MultiServerMCPClient(config)
                self._server_configs = server_configs
            except Exception as e:
                logger.error("Failed to initialize MCP client: %s", e)
                return None
        return self._client

    async def get_tools(
        self,
        server_configs: str,
    ) -> list[Callable[..., object]]:
        """Get MCP tools for a specific server, initializing client if needed.

        Args:
            server_configs: JSON string of server configurations.

        Returns:
            List of callable tools from the MCP server.
        """
        # Return cached tools if available
        if self._tools_cache:
            return self._tools_cache

        try:
            client = await self.get_client(server_configs)
            if client is None:
                return []

            # Get all tools from client
            all_tools = await client.get_tools()
            tools: list[Callable[..., object]] = list(all_tools)

            self._tools_cache = tools
            logger.info(f"Loaded {len(tools)} tools from MCP server.")
            for idx, tool in enumerate(tools):
                logger.info(
                    f"Tool {idx + 1}: {getattr(tool, 'name', 'Unknown')}, "
                    f"{getattr(tool, 'description', 'No description')}"
                )
            return tools
        except Exception as e:
            logger.warning("Failed to load tools from MCP server: %s", e)
            self._tools_cache = []
            return []

    def clear_cache(self) -> None:
        """Clear the MCP client and tools cache."""
        self._client = None
        self._tools_cache = []


# Global instance for backward compatibility
_default_adapter: MCPAdapter | None = None


def _get_default_adapter() -> MCPAdapter:
    """Get the default MCPAdapter instance."""
    global _default_adapter
    if _default_adapter is None:
        _default_adapter = MCPAdapter()
    return _default_adapter


async def get_mcp_client(
    server_configs: str,
) -> MultiServerMCPClient | None:
    """Get or initialize the global MCP client with given server configurations.

    Args:
        server_configs: JSON string of server configurations.

    Returns:
        MultiServerMCPClient instance or None if initialization fails.
    """
    return await _get_default_adapter().get_client(server_configs)


async def get_mcp_tools(
    server_configs: str,
) -> list[Callable[..., object]]:
    """Get MCP tools for a specific server, initializing client if needed.

    Args:
        server_configs: JSON string of server configurations.

    Returns:
        List of callable tools from the MCP server.
    """
    return await _get_default_adapter().get_tools(server_configs)


def clear_mcp_cache() -> None:
    """Clear the MCP client and tools cache (useful for testing)."""
    _get_default_adapter().clear_cache()


def validate_mcp_servers(
    server_configs: str | dict[str, Any],
) -> MCPValidationResult:
    """Validate MCP server configurations.

    Args:
        server_configs: Server configurations as JSON string or dictionary.

    Returns:
        MCPValidationResult with validation status and details.
    """
    if isinstance(server_configs, str):
        try:
            config = json.loads(server_configs)
        except json.JSONDecodeError as e:
            return MCPValidationResult(
                success=False,
                error=f"Invalid JSON: {e}",
            )
    else:
        config = server_configs
    return _get_default_adapter().validate(config)
