"""
Phase 107: The MCP Bridge (Pillar IV)

Translates standardized external MCP tools (like GitHub, Slack, Postgres)
into callable functions for the local Swarm models. It connects to external
MCP servers (e.g. via stdio), queries their tools, and dynamically registers
them into the Brynq OS using `dynamic_tools.register_tool`.
"""

import asyncio
import logging
from typing import Dict, Any, List

from brynq.dynamic_tools import register_tool

logger = logging.getLogger("brynq.mcp_adapter")

# Global registry of active MCP client sessions
_active_clients = {}

async def connect_external_mcp_server(name: str, command: str, args: List[str], env: Dict[str, str] = None):
    """
    Connect to an external standard MCP server via stdio.
    E.g. command="npx", args=["-y", "@modelcontextprotocol/server-github"]
    """
    try:
        from mcp import ClientSession, StdioServerParameters
        from mcp.client.stdio import stdio_client
    except ImportError:
        logger.error("The 'mcp' python package is required for the external adapter. (pip install mcp)")
        return False
        
    server_params = StdioServerParameters(
        command=command,
        args=args,
        env=env
    )
    
    try:
        # Keep references alive in the global registry
        ctx = stdio_client(server_params)
        read_stream, write_stream = await ctx.__aenter__()
        
        session = ClientSession(read_stream, write_stream)
        await session.__aenter__()
        await session.initialize()
        
        _active_clients[name] = {
            "ctx": ctx,
            "session": session,
            "command": command
        }
        
        logger.info(f"MCP Adapter: Connected to external server '{name}'")
        
        # Discover and bridge tools
        await _bridge_tools(name, session)
        return True
    except Exception as e:
        logger.error(f"MCP Adapter failed to connect to {name}: {e}")
        return False

async def _bridge_tools(server_name: str, session: Any):
    """Query the external MCP server for tools and register them natively in Brynq."""
    try:
        response = await session.list_tools()
        tools = response.tools
        registered = 0
        
        for tool in tools:
            # Create a localized closure to execute this specific tool on the remote server
            def make_handler(tool_name: str):
                def _handler(args: dict) -> dict:
                    # The handler is synchronous, but the MCP client is async.
                    # We must run it safely in the event loop.
                    try:
                        loop = asyncio.get_event_loop()
                    except RuntimeError:
                        loop = asyncio.new_event_loop()
                        asyncio.set_event_loop(loop)
                        
                    async def _call():
                        try:
                            # Note: In mcp SDK, it returns CallToolResult
                            res = await session.call_tool(tool_name, arguments=args)
                            return {"status": "success", "content": [c.model_dump() for c in res.content]}
                        except Exception as e:
                            return {"status": "error", "error": str(e)}
                            
                    if loop.is_running():
                        # If called from an already running loop (like an API route)
                        import nest_asyncio
                        nest_asyncio.apply()
                        return loop.run_until_complete(_call())
                    else:
                        return loop.run_until_complete(_call())
            
            # Namespace the tool to prevent collisions (e.g., github_search_repositories)
            brynq_tool_name = f"{server_name}_{tool.name}"
            
            register_tool(
                name=brynq_tool_name,
                description=f"[Via {server_name} MCP] {tool.description or ''}",
                handler_fn=make_handler(tool.name),
                parameters=tool.inputSchema,
                category=f"mcp_{server_name}"
            )
            registered += 1
            
        logger.info(f"MCP Adapter: Bridged {registered} tools from {server_name} into the Brynq OS.")
    except Exception as e:
        logger.error(f"MCP Adapter: Failed to bridge tools from {server_name}: {e}")

async def close_all():
    """Cleanly close all external MCP connections."""
    for name, client in _active_clients.items():
        try:
            await client["session"].__aexit__(None, None, None)
            await client["ctx"].__aexit__(None, None, None)
        except:
            pass
    _active_clients.clear()
