"""
MCP Server Registry - Manage multiple MCP server configurations

Provides:
- Persistent storage of MCP server configurations
- Connection management for multiple servers
- Tool aggregation from all connected servers
"""

import asyncio
import json
from pathlib import Path
from typing import Any, Optional

from groknroll.mcp.client import MCPClient
from groknroll.mcp.types import (
    MCPServerConfig,
    MCPServerInfo,
    MCPTool,
    MCPToolResult,
)


class MCPRegistry:
    """
    Registry for managing MCP server configurations and connections

    Stores configurations in ~/.config/groknroll/mcp.json and project-level
    .groknroll/mcp.json files.

    Example:
        registry = MCPRegistry()

        # Add a server
        config = MCPServerConfig(
            name="filesystem",
            command="npx",
            args=["-y", "@modelcontextprotocol/server-filesystem", "/path"]
        )
        registry.add_server(config)

        # Connect and use
        await registry.connect_all()
        tools = registry.list_all_tools()
        result = await registry.call_tool("filesystem", "read_file", path="/etc/hosts")
    """

    # Config file locations
    GLOBAL_CONFIG = Path.home() / ".config" / "groknroll" / "mcp.json"
    PROJECT_CONFIG = ".groknroll/mcp.json"

    def __init__(self, project_path: Optional[Path] = None):
        """
        Initialize MCP registry

        Args:
            project_path: Project directory (for project-level configs)
        """
        self.project_path = project_path or Path.cwd()
        self._configs: dict[str, MCPServerConfig] = {}
        self._clients: dict[str, MCPClient] = {}
        self._tools_cache: dict[str, list[MCPTool]] = {}

        # Load configs
        self._load_configs()

    def _load_configs(self) -> None:
        """Load configurations from global and project config files"""
        # Load global config
        if self.GLOBAL_CONFIG.exists():
            self._load_config_file(self.GLOBAL_CONFIG)

        # Load project config (overrides global)
        project_config = self.project_path / self.PROJECT_CONFIG
        if project_config.exists():
            self._load_config_file(project_config)

    def _load_config_file(self, path: Path) -> None:
        """Load configurations from a specific file"""
        try:
            with open(path) as f:
                data = json.load(f)

            for server_data in data.get("mcpServers", []):
                config = MCPServerConfig.from_dict(server_data)
                self._configs[config.name] = config
        except (json.JSONDecodeError, KeyError, ValueError):
            # Invalid config file - skip
            pass

    def _save_configs(self, global_config: bool = True) -> None:
        """Save configurations to config file"""
        path = self.GLOBAL_CONFIG if global_config else self.project_path / self.PROJECT_CONFIG

        # Ensure directory exists
        path.parent.mkdir(parents=True, exist_ok=True)

        # Build config data
        data = {"mcpServers": [config.to_dict() for config in self._configs.values()]}

        # Write to file
        with open(path, "w") as f:
            json.dump(data, f, indent=2)

    def add_server(
        self,
        config: MCPServerConfig,
        save: bool = True,
        global_config: bool = True,
    ) -> None:
        """
        Add an MCP server configuration

        Args:
            config: Server configuration
            save: Whether to persist the configuration
            global_config: Save to global config (vs project config)
        """
        self._configs[config.name] = config

        if save:
            self._save_configs(global_config)

    def remove_server(
        self,
        name: str,
        save: bool = True,
        global_config: bool = True,
    ) -> None:
        """
        Remove an MCP server configuration

        Args:
            name: Server name
            save: Whether to persist the change
            global_config: Save to global config (vs project config)
        """
        if name not in self._configs:
            raise KeyError(f"Server '{name}' not found")

        # Disconnect if connected
        if name in self._clients:
            asyncio.create_task(self._clients[name].disconnect())
            del self._clients[name]
            self._tools_cache.pop(name, None)

        del self._configs[name]

        if save:
            self._save_configs(global_config)

    def get_server(self, name: str) -> Optional[MCPServerConfig]:
        """Get server configuration by name"""
        return self._configs.get(name)

    def list_servers(self) -> list[MCPServerConfig]:
        """List all server configurations"""
        return list(self._configs.values())

    async def connect(self, name: str) -> MCPServerInfo:
        """
        Connect to a specific MCP server

        Args:
            name: Server name

        Returns:
            Server information

        Raises:
            KeyError: If server not found
            RuntimeError: If connection fails
        """
        config = self._configs.get(name)
        if not config:
            raise KeyError(f"Server '{name}' not found")

        if not config.enabled:
            raise RuntimeError(f"Server '{name}' is disabled")

        # Disconnect existing client if any
        if name in self._clients:
            await self._clients[name].disconnect()

        # Create and connect new client
        client = MCPClient(config)
        server_info = await client.connect()

        self._clients[name] = client

        # Cache tools
        tools = await client.list_tools()
        self._tools_cache[name] = tools

        return server_info

    async def disconnect(self, name: str) -> None:
        """
        Disconnect from a specific MCP server

        Args:
            name: Server name
        """
        if name in self._clients:
            await self._clients[name].disconnect()
            del self._clients[name]
            self._tools_cache.pop(name, None)

    async def connect_all(self) -> dict[str, MCPServerInfo]:
        """
        Connect to all enabled MCP servers

        Returns:
            Dictionary of server name to server info
        """
        results = {}

        for name, config in self._configs.items():
            if config.enabled:
                try:
                    info = await self.connect(name)
                    results[name] = info
                except Exception:
                    # Log error but continue with other servers
                    pass

        return results

    async def disconnect_all(self) -> None:
        """Disconnect from all MCP servers"""
        for name in list(self._clients.keys()):
            await self.disconnect(name)

    def is_connected(self, name: str) -> bool:
        """Check if a server is connected"""
        return name in self._clients and self._clients[name].is_connected

    def list_connected_servers(self) -> list[str]:
        """List names of connected servers"""
        return [name for name, client in self._clients.items() if client.is_connected]

    def list_all_tools(self) -> dict[str, list[MCPTool]]:
        """
        List all tools from all connected servers

        Returns:
            Dictionary of server name to list of tools
        """
        return dict(self._tools_cache)

    def get_tools(self, server: str) -> list[MCPTool]:
        """
        Get tools from a specific server

        Args:
            server: Server name

        Returns:
            List of tools
        """
        return self._tools_cache.get(server, [])

    async def call_tool(
        self,
        server: str,
        tool: str,
        **arguments: Any,
    ) -> MCPToolResult:
        """
        Call a tool on a specific server

        Args:
            server: Server name
            tool: Tool name
            **arguments: Tool arguments

        Returns:
            Tool execution result

        Raises:
            KeyError: If server not found
            RuntimeError: If server not connected
        """
        if server not in self._clients:
            raise KeyError(f"Server '{server}' not connected")

        client = self._clients[server]
        return await client.call_tool(tool, **arguments)

    def find_tool(self, tool_name: str) -> Optional[tuple[str, MCPTool]]:
        """
        Find a tool by name across all servers

        Args:
            tool_name: Tool name to find

        Returns:
            Tuple of (server_name, tool) or None if not found
        """
        for server_name, tools in self._tools_cache.items():
            for tool in tools:
                if tool.name == tool_name:
                    return (server_name, tool)
        return None

    async def call_tool_by_name(
        self,
        tool_name: str,
        **arguments: Any,
    ) -> MCPToolResult:
        """
        Call a tool by name (finds the server automatically)

        Args:
            tool_name: Tool name
            **arguments: Tool arguments

        Returns:
            Tool execution result

        Raises:
            KeyError: If tool not found
        """
        result = self.find_tool(tool_name)
        if not result:
            raise KeyError(f"Tool '{tool_name}' not found in any connected server")

        server_name, tool = result
        return await self.call_tool(server_name, tool_name, **arguments)

    def enable_server(self, name: str, save: bool = True) -> None:
        """Enable a server"""
        if name not in self._configs:
            raise KeyError(f"Server '{name}' not found")

        self._configs[name].enabled = True

        if save:
            self._save_configs()

    def disable_server(self, name: str, save: bool = True) -> None:
        """Disable a server"""
        if name not in self._configs:
            raise KeyError(f"Server '{name}' not found")

        self._configs[name].enabled = False

        # Disconnect if connected
        if name in self._clients:
            asyncio.create_task(self.disconnect(name))

        if save:
            self._save_configs()


# Singleton instance
_registry: Optional[MCPRegistry] = None


def get_registry(project_path: Optional[Path] = None) -> MCPRegistry:
    """
    Get the global MCP registry instance

    Args:
        project_path: Project directory (only used on first call)

    Returns:
        MCPRegistry singleton instance
    """
    global _registry
    if _registry is None:
        _registry = MCPRegistry(project_path)
    return _registry


def reset_registry() -> None:
    """Reset the global MCP registry (for testing)"""
    global _registry
    _registry = None
