# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
MCP Manager

Manages multiple MCP server connections and provides a unified interface
for tool discovery and execution across all connected servers.
"""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Callable, Dict, List, Optional

from .client import MCPClient, MCPConnectionError, MCPTool, MCPToolResult
from .config import MCPConfig, MCPServerConfig

logger = logging.getLogger(__name__)


class ServerStatus(str, Enum):
    """Status of an MCP server connection."""

    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    ERROR = "error"
    DISABLED = "disabled"


@dataclass
class ServerState:
    """Runtime state for an MCP server."""

    config: MCPServerConfig
    client: Optional[MCPClient] = None
    status: ServerStatus = ServerStatus.DISCONNECTED
    error_message: Optional[str] = None
    connected_at: Optional[datetime] = None
    last_error_at: Optional[datetime] = None
    reconnect_attempts: int = 0
    tools: List[MCPTool] = field(default_factory=list)


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

    Provides:
    - Concurrent connection management
    - Automatic reconnection on failure
    - Unified tool interface across servers
    - Server health monitoring

    Usage:
        config = MCPConfig(servers=[...])
        manager = MCPManager(config)

        # Connect to all servers
        await manager.connect_all()

        # Get all tools from all servers
        tools = await manager.list_all_tools()

        # Call a tool (manager routes to correct server)
        result = await manager.call_tool("github:create_issue", {"title": "Bug"})

        # Disconnect all
        await manager.disconnect_all()
    """

    def __init__(self, config: MCPConfig):
        self.config = config
        self._servers: Dict[str, ServerState] = {}
        self._lock = asyncio.Lock()
        self._reconnect_tasks: Dict[str, asyncio.Task] = {}

        # Event callbacks
        self._on_connect: List[Callable[[str], None]] = []
        self._on_disconnect: List[Callable[[str, Optional[str]], None]] = []
        self._on_tools_changed: List[Callable[[str, List[MCPTool]], None]] = []

        # Initialize server states
        for server_config in config.servers:
            self._servers[server_config.name] = ServerState(config=server_config)

    # ----------------------------------------------------------
    # Connection Management
    # ----------------------------------------------------------

    async def connect_all(self) -> Dict[str, bool]:
        """
        Connect to all enabled servers concurrently.

        Returns:
            Dict mapping server names to connection success status
        """
        results = {}
        tasks = []

        for name, state in self._servers.items():
            if state.config.enabled and state.config.auto_connect:
                tasks.append(self._connect_server(name))

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

            i = 0
            for name, state in self._servers.items():
                if state.config.enabled and state.config.auto_connect:
                    if isinstance(outcomes[i], Exception):
                        results[name] = False
                        logger.error(f"Failed to connect to '{name}': {outcomes[i]}")
                    else:
                        results[name] = outcomes[i]
                    i += 1

        connected = sum(1 for v in results.values() if v)
        logger.info(f"Connected to {connected}/{len(results)} MCP servers")

        return results

    async def disconnect_all(self) -> None:
        """Disconnect from all servers."""
        # Cancel any reconnect tasks
        for task in self._reconnect_tasks.values():
            task.cancel()
        self._reconnect_tasks.clear()

        # Disconnect all servers
        tasks = [self._disconnect_server(name) for name in self._servers]

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

        logger.info("Disconnected from all MCP servers")

    async def connect_server(self, name: str) -> bool:
        """Connect to a specific server by name."""
        if name not in self._servers:
            raise ValueError(f"Unknown server: {name}")
        return await self._connect_server(name)

    async def disconnect_server(self, name: str) -> None:
        """Disconnect from a specific server."""
        if name not in self._servers:
            raise ValueError(f"Unknown server: {name}")
        await self._disconnect_server(name)

    async def _connect_server(self, name: str) -> bool:
        """Internal method to connect to a server."""
        state = self._servers[name]

        if state.status == ServerStatus.CONNECTED:
            return True

        async with self._lock:
            state.status = ServerStatus.CONNECTING
            state.error_message = None

            try:
                # Create client
                client = MCPClient(state.config)
                await asyncio.wait_for(client.connect(), timeout=self.config.connect_timeout)

                # Fetch tools
                tools = await client.list_tools()

                # Update state
                state.client = client
                state.status = ServerStatus.CONNECTED
                state.connected_at = datetime.now(timezone.utc)
                state.reconnect_attempts = 0
                state.tools = tools

                # Notify listeners
                for callback in self._on_connect:
                    try:
                        callback(name)
                    except Exception as e:
                        logger.error(f"Error in connect callback: {e}")

                for callback in self._on_tools_changed:
                    try:
                        callback(name, tools)
                    except Exception as e:
                        logger.error(f"Error in tools_changed callback: {e}")

                logger.info(f"Connected to MCP server '{name}' with {len(tools)} tools")
                return True

            except asyncio.TimeoutError:
                state.status = ServerStatus.ERROR
                state.error_message = "Connection timeout"
                state.last_error_at = datetime.now(timezone.utc)
                logger.error(f"Timeout connecting to MCP server '{name}'")

                if self.config.reconnect_on_failure:
                    self._schedule_reconnect(name)

                return False

            except Exception as e:
                state.status = ServerStatus.ERROR
                state.error_message = str(e)
                state.last_error_at = datetime.now(timezone.utc)
                logger.error(f"Failed to connect to MCP server '{name}': {e}")

                if self.config.reconnect_on_failure:
                    self._schedule_reconnect(name)

                return False

    async def _disconnect_server(self, name: str) -> None:
        """Internal method to disconnect from a server."""
        state = self._servers[name]

        # Cancel reconnect task if any
        if name in self._reconnect_tasks:
            self._reconnect_tasks[name].cancel()
            del self._reconnect_tasks[name]

        if state.client:
            try:
                await state.client.disconnect()
            except Exception as e:
                logger.warning(f"Error disconnecting from '{name}': {e}")

        old_status = state.status
        state.client = None
        state.status = ServerStatus.DISCONNECTED
        state.tools = []

        if old_status == ServerStatus.CONNECTED:
            for callback in self._on_disconnect:
                try:
                    callback(name, None)
                except Exception as e:
                    logger.error(f"Error in disconnect callback: {e}")

    def _schedule_reconnect(self, name: str) -> None:
        """Schedule a reconnection attempt."""
        state = self._servers[name]

        if state.reconnect_attempts >= self.config.max_reconnect_attempts:
            logger.warning(
                f"Max reconnect attempts ({self.config.max_reconnect_attempts}) "
                f"reached for server '{name}'"
            )
            return

        state.reconnect_attempts += 1
        delay = self.config.reconnect_delay * state.reconnect_attempts

        async def reconnect():
            await asyncio.sleep(delay)
            logger.info(f"Attempting reconnection to '{name}' (attempt {state.reconnect_attempts})")
            await self._connect_server(name)

        self._reconnect_tasks[name] = asyncio.create_task(reconnect())

    # ----------------------------------------------------------
    # Server Information
    # ----------------------------------------------------------

    def get_server_status(self, name: str) -> ServerStatus:
        """Get the status of a server."""
        if name not in self._servers:
            raise ValueError(f"Unknown server: {name}")
        return self._servers[name].status

    def get_connected_servers(self) -> List[str]:
        """Get names of all connected servers."""
        return [
            name for name, state in self._servers.items() if state.status == ServerStatus.CONNECTED
        ]

    def get_server_info(self, name: str) -> Dict[str, Any]:
        """Get detailed info about a server."""
        if name not in self._servers:
            raise ValueError(f"Unknown server: {name}")

        state = self._servers[name]

        return {
            "name": name,
            "status": state.status.value,
            "enabled": state.config.enabled,
            "transport": state.config.transport.value,
            "error": state.error_message,
            "connected_at": state.connected_at.isoformat() if state.connected_at else None,
            "tool_count": len(state.tools),
            "server_info": state.client.server_info if state.client else None,
        }

    def get_all_server_info(self) -> Dict[str, Dict[str, Any]]:
        """Get info about all servers."""
        return {name: self.get_server_info(name) for name in self._servers}

    # ----------------------------------------------------------
    # Tool Operations
    # ----------------------------------------------------------

    async def list_all_tools(self, force_refresh: bool = False) -> List[MCPTool]:
        """
        List all tools from all connected servers.

        Returns tools with names prefixed by server name (e.g., "github:create_issue")
        """
        all_tools = []

        for name, state in self._servers.items():
            if state.status != ServerStatus.CONNECTED or not state.client:
                continue

            try:
                if force_refresh:
                    state.tools = await state.client.list_tools()

                # Prefix tool names with server name
                for tool in state.tools:
                    prefixed_tool = MCPTool(
                        name=f"{name}{self.config.tool_prefix_separator}{tool.name}",
                        description=f"[{name}] {tool.description}",
                        input_schema=tool.input_schema,
                        annotations=tool.annotations,
                    )
                    all_tools.append(prefixed_tool)

            except Exception as e:
                logger.error(f"Error listing tools from '{name}': {e}")

        return all_tools

    async def list_server_tools(self, name: str) -> List[MCPTool]:
        """List tools from a specific server."""
        if name not in self._servers:
            raise ValueError(f"Unknown server: {name}")

        state = self._servers[name]
        if state.status != ServerStatus.CONNECTED or not state.client:
            raise MCPConnectionError(f"Server '{name}' is not connected")

        return await state.client.list_tools()

    async def call_tool(
        self,
        full_name: str,
        arguments: Optional[Dict[str, Any]] = None,
    ) -> MCPToolResult:
        """
        Call a tool by its full name (server:tool_name).

        Args:
            full_name: Full tool name including server prefix
            arguments: Tool arguments

        Returns:
            MCPToolResult with content
        """
        # Parse server and tool name
        sep = self.config.tool_prefix_separator
        if sep not in full_name:
            raise ValueError(
                f"Invalid tool name '{full_name}'. Expected format: server{sep}tool_name"
            )

        server_name, tool_name = full_name.split(sep, 1)

        if server_name not in self._servers:
            raise ValueError(f"Unknown server: {server_name}")

        state = self._servers[server_name]
        if state.status != ServerStatus.CONNECTED or not state.client:
            raise MCPConnectionError(f"Server '{server_name}' is not connected")

        return await state.client.call_tool(tool_name, arguments)

    def get_tool(self, full_name: str) -> Optional[MCPTool]:
        """Get tool definition by full name."""
        sep = self.config.tool_prefix_separator
        if sep not in full_name:
            return None

        server_name, tool_name = full_name.split(sep, 1)

        state = self._servers.get(server_name)
        if not state:
            return None

        for tool in state.tools:
            if tool.name == tool_name:
                return MCPTool(
                    name=full_name,
                    description=f"[{server_name}] {tool.description}",
                    input_schema=tool.input_schema,
                    annotations=tool.annotations,
                )

        return None

    # ----------------------------------------------------------
    # Dynamic Server Management
    # ----------------------------------------------------------

    async def add_server(
        self,
        config: MCPServerConfig,
        connect: bool = True,
    ) -> bool:
        """
        Add a new server at runtime.

        Args:
            config: Server configuration
            connect: Whether to immediately connect

        Returns:
            True if successfully added (and connected if requested)
        """
        if config.name in self._servers:
            raise ValueError(f"Server '{config.name}' already exists")

        self._servers[config.name] = ServerState(config=config)
        self.config.servers.append(config)

        if connect and config.enabled:
            return await self._connect_server(config.name)

        return True

    async def remove_server(self, name: str) -> bool:
        """
        Remove a server at runtime.

        Args:
            name: Server name to remove

        Returns:
            True if successfully removed
        """
        if name not in self._servers:
            return False

        await self._disconnect_server(name)
        del self._servers[name]

        # Remove from config
        self.config.servers = [s for s in self.config.servers if s.name != name]

        return True

    # ----------------------------------------------------------
    # Event Callbacks
    # ----------------------------------------------------------

    def on_connect(self, callback: Callable[[str], None]) -> None:
        """Register callback for server connect events."""
        self._on_connect.append(callback)

    def on_disconnect(self, callback: Callable[[str, Optional[str]], None]) -> None:
        """Register callback for server disconnect events."""
        self._on_disconnect.append(callback)

    def on_tools_changed(self, callback: Callable[[str, List[MCPTool]], None]) -> None:
        """Register callback for tools changed events."""
        self._on_tools_changed.append(callback)

    # ----------------------------------------------------------
    # Context Manager
    # ----------------------------------------------------------

    async def __aenter__(self) -> "MCPManager":
        """Async context manager entry."""
        await self.connect_all()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit."""
        await self.disconnect_all()
