"""
MCP Client utilities for connecting to MCP servers.

This module provides a wrapper around the MCP client SDK for connecting
to and calling tools on MCP servers.
"""

import asyncio
from contextlib import AsyncExitStack, asynccontextmanager
from pathlib import Path
from typing import Any, Dict, List, Optional

from tenacity import (
    AsyncRetrying,
    retry_if_exception_message,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from loguru import logger

# Allowed commands for MCP server execution (security whitelist)
ALLOWED_COMMANDS = {"node", "npx", "python", "python3", "uv", "uvx", "docker"}

# Retry configuration for subprocess startup race condition
INIT_RETRY_ATTEMPTS = 5
INIT_RETRY_BASE_DELAY = 0.1  # seconds
INIT_RETRY_MAX_DELAY = 2.0  # seconds
INIT_RETRY_BACKOFF_FACTOR = 2.0

# The MCP SDK is an optional dependency (installed via the [mcp] extras group),
# so we guard the import and set a flag for runtime availability checks.
try:
    from mcp import ClientSession
    from mcp.client.stdio import StdioServerParameters, stdio_client

    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False
    logger.warning(
        "MCP client not available. Install with: pip install mcp[cli]"
    )


class MCPClientError(Exception):
    """Error during MCP client operations."""

    pass


class MCPClient:
    """
    Client for connecting to and calling tools on MCP servers.

    Usage:
        async with MCPClient(config) as client:
            tools = await client.list_tools()
            result = await client.call_tool("tool_name", {"arg": "value"})
    """

    def __init__(
        self,
        server_config: Dict[str, Any],
        timeout: float = 60.0,
    ):
        """
        Initialize MCP client.

        Args:
            server_config: Server configuration with keys:
                - command: str - Command to run (e.g., "python")
                - args: List[str] - Command arguments
                - env: Dict[str, str] - Environment variables (optional)
                - name: str - Server name for logging (optional)
            timeout: Timeout in seconds for tool calls
        """
        if not MCP_AVAILABLE:
            raise MCPClientError(
                "MCP client not available. Install with: pip install mcp[cli]"
            )

        # Validate server configuration for security
        self._validate_server_config(server_config)

        self.config = server_config
        self.timeout = timeout
        self.name = server_config.get("name", "unnamed")
        self._session: Optional[ClientSession] = None
        self._read_stream = None
        self._write_stream = None
        self._connected = False

    def _validate_server_config(self, config: Dict[str, Any]) -> None:
        """
        Validate server configuration for security.

        Args:
            config: Server configuration dictionary

        Raises:
            MCPClientError: If configuration is invalid or insecure
        """
        # Check command exists
        command = config.get("command", "")
        command = command.strip()
        if not command:
            raise MCPClientError("Server config missing 'command'")

        if " " in command or "\t" in command:
            raise MCPClientError(
                "'command' must be a single executable name without spaces; "
                "pass arguments via the 'args' field"
            )

        # Write stripped command back for use by connect()
        config["command"] = command

        # Extract base command (handle paths like /usr/bin/node)
        base_cmd = Path(command).name
        if base_cmd not in ALLOWED_COMMANDS:
            raise MCPClientError(
                f"Command '{base_cmd}' not in allowed list: {ALLOWED_COMMANDS}. "
                f"Add to ALLOWED_COMMANDS if this is a trusted command."
            )

        # Validate args are a list of strings
        args = config.get("args", [])
        if not isinstance(args, list):
            raise MCPClientError("Server 'args' must be a list")
        if not all(isinstance(a, str) for a in args):
            raise MCPClientError("Server 'args' must contain only strings")

        # Validate env is a dict of strings if provided
        env = config.get("env")
        if env is not None:
            if not isinstance(env, dict):
                raise MCPClientError("Server 'env' must be a dictionary")
            if not all(
                isinstance(k, str) and isinstance(v, str)
                for k, v in env.items()
            ):
                raise MCPClientError(
                    "Server 'env' must contain only string keys and values"
                )

    @asynccontextmanager
    async def connect(self):
        """
        Connect to the MCP server as an async context manager.

        Yields:
            self: The connected client instance
        """
        server_params = StdioServerParameters(
            command=self.config["command"],
            args=self.config.get("args", []),
            env=self.config.get("env"),
        )

        logger.info(f"Connecting to MCP server '{self.name}'...")

        try:
            async with stdio_client(server_params) as (
                read_stream,
                write_stream,
            ):
                self._read_stream = read_stream
                self._write_stream = write_stream

                # ClientSession must be used as async context manager
                async with ClientSession(read_stream, write_stream) as session:
                    self._session = session

                    await self._initialize_with_retry()

                    self._connected = True
                    logger.info(f"Connected to MCP server '{self.name}'")

                    yield self

        except Exception as e:
            logger.exception(f"Failed to connect to MCP server '{self.name}'")
            raise MCPClientError(f"Connection failed: {e}") from e
        finally:
            self._connected = False
            self._session = None
            self._read_stream = None
            self._write_stream = None

    async def _initialize_with_retry(self) -> None:
        """Initialize the MCP session with retry logic for subprocess startup race.

        Retries on TimeoutError and "before initialization" errors using
        exponential backoff, then wraps final failures in MCPClientError.
        """
        try:
            async for attempt in AsyncRetrying(
                stop=stop_after_attempt(INIT_RETRY_ATTEMPTS),
                wait=wait_exponential(
                    multiplier=INIT_RETRY_BASE_DELAY,
                    max=INIT_RETRY_MAX_DELAY,
                    exp_base=INIT_RETRY_BACKOFF_FACTOR,
                ),
                retry=(
                    retry_if_exception_type(asyncio.TimeoutError)
                    | retry_if_exception_message(
                        match="(?i)before initialization"
                    )
                ),
                before_sleep=lambda rs: logger.debug(
                    f"Server '{self.name}' not ready "
                    f"(attempt {rs.attempt_number}/{INIT_RETRY_ATTEMPTS}), retrying..."
                ),
                reraise=True,
            ):
                with attempt:
                    if self._session is None:
                        raise MCPClientError("Session not initialized")  # noqa: TRY301 — re-raised by except MCPClientError
                    await asyncio.wait_for(
                        self._session.initialize(),
                        timeout=min(5.0, self.timeout),
                    )
        except asyncio.TimeoutError as e:
            raise MCPClientError(
                f"Timeout initializing connection to '{self.name}' "
                f"after {INIT_RETRY_ATTEMPTS} attempts"
            ) from e
        except MCPClientError:
            raise
        except Exception as e:
            if "before initialization" in str(e).lower():
                raise MCPClientError(
                    f"Failed to initialize connection to '{self.name}' "
                    f"after {INIT_RETRY_ATTEMPTS} attempts: {e}"
                ) from e
            raise

    async def list_tools(self) -> List[Dict[str, Any]]:
        """
        List available tools on the connected server.

        Returns:
            List of tool definitions with name, description, and input schema
        """
        if not self._connected or not self._session:
            raise MCPClientError("Not connected to server")

        try:
            result = await asyncio.wait_for(
                self._session.list_tools(),
                timeout=self.timeout,
            )
            tools = []
            for tool in result.tools:
                tools.append(
                    {
                        "name": tool.name,
                        "description": tool.description or "",
                        "input_schema": (
                            tool.inputSchema
                            if hasattr(tool, "inputSchema")
                            else {}
                        ),
                    }
                )
            return tools
        except asyncio.TimeoutError:
            logger.warning(
                f"Timeout listing tools from '{self.name}' after {self.timeout}s"
            )
            raise MCPClientError(f"Timeout listing tools after {self.timeout}s")
        except Exception as e:
            logger.exception(f"Failed to list tools from '{self.name}'")
            raise MCPClientError(f"Failed to list tools: {e}") from e

    async def call_tool(
        self,
        name: str,
        arguments: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Call a tool on the connected server.

        Args:
            name: Name of the tool to call
            arguments: Arguments to pass to the tool

        Returns:
            Tool result as a dictionary
        """
        if not self._connected or not self._session:
            raise MCPClientError("Not connected to server")

        try:
            logger.debug(
                f"Calling tool '{name}' on '{self.name}' with args: {arguments}"
            )

            result = await asyncio.wait_for(
                self._session.call_tool(name, arguments or {}),
                timeout=self.timeout,
            )

            # Parse the result
            if result.isError:
                return {
                    "status": "error",
                    "error": str(result.content),
                }

            # Extract content from the result
            content = []
            # Check if result.content is iterable (not None or other non-iterable types)
            if result.content is not None and hasattr(
                result.content, "__iter__"
            ):
                for item in result.content:
                    if hasattr(item, "text"):
                        content.append(item.text)
                    elif hasattr(item, "data"):
                        content.append(str(item.data))
                    else:
                        content.append(str(item))
            elif result.content is not None:
                # Handle non-iterable content by converting to string
                content.append(str(result.content))

            return {
                "status": "success",
                "content": "\n".join(content) if content else "",
                "raw": result.content,
            }

        except asyncio.TimeoutError:
            logger.exception(
                f"Tool call '{name}' timed out after {self.timeout}s"
            )
            raise MCPClientError(f"Tool call timed out after {self.timeout}s")
        except Exception as e:
            logger.exception(f"Failed to call tool '{name}' on '{self.name}'")
            raise MCPClientError(f"Tool call failed: {e}") from e


class MCPClientManager:
    """
    Manager for multiple MCP client connections.

    Handles connecting to multiple MCP servers and aggregating their tools.
    """

    def __init__(self, server_configs: List[Dict[str, Any]]):
        """
        Initialize the manager with server configurations.

        Args:
            server_configs: List of server configurations
        """
        self.server_configs = server_configs
        self._clients: Dict[str, MCPClient] = {}

    @asynccontextmanager
    async def connect_all(self):
        """
        Connect to all configured MCP servers.

        Yields:
            self: The manager with all clients connected
        """
        # Create clients for each server
        clients = [MCPClient(config) for config in self.server_configs]

        # Connect to all servers
        # Connect sequentially to avoid overwhelming the system.
        # AsyncExitStack manages cleanup of successfully-entered contexts.
        async with AsyncExitStack() as stack:
            for client in clients:
                ctx = None
                try:
                    ctx = client.connect()
                    connected = await stack.enter_async_context(ctx)
                    self._clients[client.name] = connected
                except Exception:
                    logger.warning(
                        f"Failed to connect to server '{client.name}'. Skipping."
                    )
                    # If ctx was created but __aenter__ failed, clean up
                    # manually — AsyncExitStack only tracks contexts that
                    # were successfully entered, but the MCP subprocess
                    # may already be running.
                    if ctx is not None:
                        try:
                            await ctx.__aexit__(None, None, None)
                        except Exception:
                            logger.debug(
                                "best-effort cleanup of partially-entered async context",
                                exc_info=True,
                            )

            try:
                yield self
            finally:
                self._clients.clear()

    async def list_all_tools(self) -> Dict[str, List[Dict[str, Any]]]:
        """
        List tools from all connected servers.

        Returns:
            Dictionary mapping server name to list of tools
        """
        all_tools = {}
        for name, client in self._clients.items():
            try:
                tools = await client.list_tools()
                all_tools[name] = tools
            except MCPClientError:
                logger.warning(f"Failed to list tools from '{name}'")
                all_tools[name] = []
        return all_tools

    async def call_tool(
        self,
        server_name: str,
        tool_name: str,
        arguments: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Call a tool on a specific server.

        Args:
            server_name: Name of the server
            tool_name: Name of the tool to call
            arguments: Arguments to pass to the tool

        Returns:
            Tool result
        """
        if server_name not in self._clients:
            raise MCPClientError(f"Server '{server_name}' not connected")

        return await self._clients[server_name].call_tool(tool_name, arguments)

    def get_connected_servers(self) -> List[str]:
        """Get list of connected server names."""
        return list(self._clients.keys())


def run_async(coro, timeout: float = 300.0):
    """
    Run an async coroutine synchronously.

    Helper for running async MCP operations from sync code.

    Args:
        coro: The coroutine to run
        timeout: Maximum time to wait in seconds (default 5 minutes)

    Returns:
        The result of the coroutine

    Raises:
        MCPClientError: If the operation times out
    """
    import concurrent.futures

    try:
        # Check if we're already in an async context
        asyncio.get_running_loop()
        # We are in an async context - use thread pool to avoid nesting
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(
                asyncio.run, asyncio.wait_for(coro, timeout=timeout)
            )
            try:
                return future.result(timeout=timeout + 5.0)
            except TimeoutError:
                raise MCPClientError(
                    f"Async operation timed out after {timeout}s"
                )
    except RuntimeError:
        # No running event loop - safe to use asyncio.run()
        try:
            return asyncio.run(asyncio.wait_for(coro, timeout=timeout))
        except TimeoutError:
            raise MCPClientError(f"Async operation timed out after {timeout}s")
