"""MCP Client Manager for Fluxibly.

Handles MCP server lifecycle, tool discovery, and invocation.
Refactored from fluxibly.mcp_client.manager for the new tools layer.
"""

from __future__ import annotations

import os
from contextlib import AsyncExitStack
from pathlib import Path
from typing import Any

from fluxibly.exceptions import (
    MCPConnectionError,
    MCPError,
    MCPToolNotFoundError,
)
from fluxibly.logging import Logger

try:
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client
except ImportError as e:
    raise ImportError(
        "MCP package is required for MCPClientManager. "
        "Install it with: pip install fluxibly[mcp]"
    ) from e


class MCPClientManager:
    """Central manager for MCP server communication.

    Responsibilities:
    - Server lifecycle management (connect / disconnect)
    - Tool discovery and aggregation from all connected servers
    - Unified tool invocation interface
    - Configuration loading from dict or YAML
    """

    def __init__(self) -> None:
        self.servers: dict[str, dict[str, Any]] = {}
        self.tools: dict[str, dict[str, Any]] = (
            {}
        )  # tool_name -> {name, description, schema, server}
        self.logger = Logger(component="mcp")

    # ── Server Management ─────────────────────────────────────────

    async def connect_server(self, name: str, config: dict[str, Any]) -> None:
        """Connect to an MCP server and discover its tools.

        Args:
            name: Logical server name.
            config: Server config with keys: command, args, env, etc.

        Raises:
            MCPConnectionError: If connection fails.
        """
        command = config.get("command")
        args = config.get("args", [])
        env = config.get("env", {})

        if not command:
            raise MCPConnectionError(
                f"Server '{name}' missing required 'command' field"
            )

        # Expand ${VAR} references in env values
        expanded_env = {}
        for key, value in env.items():
            if (
                isinstance(value, str)
                and value.startswith("${")
                and value.endswith("}")
            ):
                env_var = value[2:-1]
                expanded_env[key] = os.environ.get(env_var, "")
            else:
                expanded_env[key] = str(value)

        # Merge with current process environment
        server_env = os.environ.copy()
        server_env.update(expanded_env)

        server_params = StdioServerParameters(
            command=command, args=args, env=server_env
        )

        stack = AsyncExitStack()
        try:
            # Use AsyncExitStack for proper resource cleanup
            read, write = await stack.enter_async_context(
                stdio_client(server_params)
            )
            session = await stack.enter_async_context(
                ClientSession(read, write)
            )
            await session.initialize()

            self.servers[name] = {
                "config": config,
                "session": session,
                "stack": stack,
            }

            # Discover tools from the newly connected server
            await self._discover_server_tools(name, session)

            self.logger.info(
                f"Connected to MCP server '{name}' — {len(self._server_tool_names(name))} tools"
            )

        except Exception as e:
            await stack.aclose()
            raise MCPConnectionError(
                f"Failed to connect to MCP server '{name}': {e}"
            ) from e

    async def disconnect_server(self, name: str) -> None:
        """Disconnect a single MCP server."""
        server_info = self.servers.pop(name, None)
        if not server_info:
            return

        try:
            stack = server_info.get("stack")
            if stack:
                await stack.aclose()
        except Exception as e:
            self.logger.warning(f"Error disconnecting server '{name}': {e}")

        # Remove tools from this server
        self.tools = {
            k: v for k, v in self.tools.items() if v.get("server") != name
        }

    # ── Batch Initialize / Shutdown ───────────────────────────────

    async def initialize(
        self, servers_config: dict[str, dict[str, Any]] | None = None
    ) -> None:
        """Connect to multiple MCP servers from a config dict.

        Args:
            servers_config: {server_name: {command, args, env, enabled, ...}}
                            If None, does nothing (servers can be added individually).
        """
        if not servers_config:
            return

        for name, cfg in servers_config.items():
            if not cfg.get("enabled", True):
                continue
            try:
                await self.connect_server(name, cfg)
            except MCPConnectionError as e:
                self.logger.error(f"Skipping MCP server '{name}': {e}")

    async def shutdown(self) -> None:
        """Gracefully disconnect all MCP servers."""
        for name in list(self.servers.keys()):
            await self.disconnect_server(name)
        self.logger.info("MCP manager shutdown complete")

    # ── Tool Discovery ────────────────────────────────────────────

    async def _discover_server_tools(
        self, server_name: str, session: ClientSession
    ) -> None:
        """Discover and register tools from a connected server."""
        try:
            tools_result = await session.list_tools()
            for tool in tools_result.tools:
                if tool.name in self.tools:
                    existing = self.tools[tool.name]["server"]
                    self.logger.warning(
                        f"Tool name conflict: '{tool.name}' in both '{existing}' and '{server_name}'. "
                        f"Keeping '{existing}'."
                    )
                    continue

                self.tools[tool.name] = {
                    "name": tool.name,
                    "description": tool.description or "",
                    "schema": tool.inputSchema,
                    "server": server_name,
                }
        except Exception as e:
            self.logger.error(
                f"Failed to discover tools from server '{server_name}': {e}"
            )

    def get_server_tools(self, server_name: str) -> list[dict[str, Any]]:
        """Get tools from a specific server.

        Returns:
            List of tool dicts: [{name, description, schema, server}, ...]
        """
        return [t for t in self.tools.values() if t["server"] == server_name]

    def get_all_tools(self) -> list[dict[str, Any]]:
        """Get all discovered tools across all servers."""
        return list(self.tools.values())

    # ── Tool Invocation ───────────────────────────────────────────

    async def invoke_tool(self, tool_name: str, args: dict[str, Any]) -> Any:
        """Execute a tool through its MCP server.

        Args:
            tool_name: Registered tool name.
            args: Tool arguments.

        Returns:
            Raw MCP tool result.

        Raises:
            MCPToolNotFoundError: If tool is not registered.
            MCPError: If invocation fails.
        """
        if tool_name not in self.tools:
            available = list(self.tools.keys())
            raise MCPToolNotFoundError(
                f"Tool '{tool_name}' not found. Available: {available}"
            )

        tool_info = self.tools[tool_name]
        server_name = tool_info["server"]
        server_info = self.servers.get(server_name)

        if not server_info:
            raise MCPError(f"Server '{server_name}' not connected")

        session: ClientSession = server_info["session"]

        try:
            self.logger.debug(f"Invoking '{tool_name}' on '{server_name}'")
            result = await session.call_tool(tool_name, arguments=args)
            return result
        except Exception as e:
            raise MCPError(f"Failed to invoke tool '{tool_name}': {e}") from e

    # ── Helpers ───────────────────────────────────────────────────

    def _server_tool_names(self, server_name: str) -> list[str]:
        """Return tool names belonging to a server."""
        return [
            name
            for name, info in self.tools.items()
            if info["server"] == server_name
        ]

    @staticmethod
    def load_config_from_yaml(config_path: str) -> dict[str, dict[str, Any]]:
        """Load MCP servers config from a YAML file.

        Expected format::

            mcp_servers:
              server_name:
                command: "python"
                args: ["-m", "server_module"]
                env:
                  KEY: "${ENV_VAR}"
                enabled: true

        Returns:
            The ``mcp_servers`` dict from the YAML.

        Raises:
            FileNotFoundError: If file doesn't exist.
            ValueError: If YAML is invalid.
        """
        import yaml

        path = Path(config_path)
        if not path.exists():
            raise FileNotFoundError(
                f"MCP config file not found: {config_path}"
            )

        with path.open("r") as f:
            config = yaml.safe_load(f)

        if not isinstance(config, dict):
            raise ValueError(
                f"Invalid MCP config: expected dict, got {type(config)}"
            )

        return config.get("mcp_servers", {})
