"""
MCP (Model Context Protocol) Client for federated tool ecosystems.

Enables integration with external tool providers via MCP servers,
allowing agents to discover and call remote tools seamlessly.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

try:
    import aiohttp
except ImportError:
    aiohttp = None


class MCPToolSource:
    """
    Client for MCP server tool discovery and invocation.
    
    MCP servers expose standardized APIs for:
    - Tool discovery (GET /tools)
    - Tool invocation (POST /call)
    
    Example:
        mcp = MCPToolSource("http://localhost:3000", api_key="my-key")
        tools = await mcp.discover()
        result = await mcp.call_tool("web_search", {"query": "AI"})
    """

    def __init__(self, base_url: str, api_key: Optional[str] = None):
        """
        Initialize MCP client.
        
        Args:
            base_url: Base URL of MCP server (e.g., "http://localhost:3000")
            api_key: Optional API key for authentication
        """
        if aiohttp is None:
            raise RuntimeError(
                "aiohttp is required for MCP client. "
                "Install with: pip install aiohttp"
            )
        
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key

    def _get_headers(self) -> Dict[str, str]:
        """Get HTTP headers with optional auth."""
        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers

    async def discover(self) -> List[Dict[str, Any]]:
        """
        Discover available tools from MCP server.
        
        Returns:
            List of tool metadata dicts with 'name', 'description', 'parameters'
        
        Raises:
            aiohttp.ClientError: On connection/HTTP errors
        """
        async with aiohttp.ClientSession() as session:
            headers = self._get_headers()
            url = f"{self.base_url}/tools"
            
            async with session.get(url, headers=headers) as response:
                response.raise_for_status()
                data = await response.json()
                
                # Normalize response format
                if isinstance(data, dict) and "tools" in data:
                    return data["tools"]
                elif isinstance(data, list):
                    return data
                else:
                    return []

    async def call_tool(
        self,
        tool_name: str,
        params: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        Call a remote tool via MCP server.
        
        Args:
            tool_name: Name of tool to invoke
            params: Tool parameters
        
        Returns:
            Tool execution result
        
        Raises:
            aiohttp.ClientError: On connection/HTTP errors
        """
        async with aiohttp.ClientSession() as session:
            headers = self._get_headers()
            url = f"{self.base_url}/call"
            payload = {"tool": tool_name, "params": params}
            
            async with session.post(url, json=payload, headers=headers) as response:
                response.raise_for_status()
                return await response.json()

    async def health_check(self) -> bool:
        """
        Check if MCP server is reachable.
        
        Returns:
            True if server responds to /health endpoint
        """
        try:
            async with aiohttp.ClientSession() as session:
                url = f"{self.base_url}/health"
                async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as response:
                    return response.status == 200
        except Exception:
            return False


class MCPToolRegistry:
    """
    Registry for managing multiple MCP tool sources.
    
    Allows agents to discover and call tools from multiple federated MCP servers.
    
    Example:
        registry = MCPToolRegistry()
        registry.add_source("primary", MCPToolSource("http://mcp1.example.com"))
        registry.add_source("backup", MCPToolSource("http://mcp2.example.com"))
        
        tools = await registry.discover_all()
        result = await registry.call_tool("web_search", {"query": "AI"})
    """

    def __init__(self):
        self.sources: Dict[str, MCPToolSource] = {}

    def add_source(self, name: str, source: MCPToolSource) -> None:
        """Register an MCP source."""
        self.sources[name] = source

    async def discover_all(self) -> Dict[str, List[Dict[str, Any]]]:
        """
        Discover tools from all registered sources.
        
        Returns:
            Dict mapping source name to list of tool metadata
        """
        import asyncio
        
        tasks = {
            name: source.discover()
            for name, source in self.sources.items()
        }
        
        results = await asyncio.gather(
            *tasks.values(),
            return_exceptions=True
        )
        
        discovered = {}
        for (name, _), result in zip(tasks.items(), results):
            if isinstance(result, Exception):
                discovered[name] = []
            else:
                discovered[name] = result
        
        return discovered

    async def call_tool(
        self,
        tool_name: str,
        params: Dict[str, Any],
        source_name: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Call a tool, optionally specifying which source to use.
        
        If source_name is None, tries all sources until one succeeds.
        
        Args:
            tool_name: Tool name
            params: Tool parameters
            source_name: Optional source to use
        
        Returns:
            Tool result
        
        Raises:
            RuntimeError: If no sources available or all fail
        """
        if source_name:
            if source_name not in self.sources:
                raise ValueError(f"Unknown MCP source: {source_name}")
            return await self.sources[source_name].call_tool(tool_name, params)
        
        # Try all sources
        last_error = None
        for name, source in self.sources.items():
            try:
                return await source.call_tool(tool_name, params)
            except Exception as e:
                last_error = e
                continue
        
        raise RuntimeError(
            f"Failed to call tool '{tool_name}' from any MCP source. "
            f"Last error: {last_error}"
        )
