"""MCP Auto-Registration - Make MCP tools first-class.

Host/runtime concern: given an MCP client and the process-wide tool registries,
register MCP server tools as normal tools.

This module intentionally does NOT depend on ``stdlib``.
"""

from __future__ import annotations

import logging
from typing import Any

from core.pathways.models.tool import ToolArgs, ToolOutput
from core.tools import get_tool_registry
from infrastructure.mcp.client import McpClientService

logger = logging.getLogger(__name__)


class MCPToolBridge:
    """Bridges MCP server tools into the AlbusOS tool registry.

    Replaces module-level globals with a proper stateful object.
    Each registered tool handler captures a reference to the bridge
    so the MCP client is always reachable.
    """

    def __init__(self, mcp_client: McpClientService) -> None:
        self._client = mcp_client
        self._registered: set[str] = set()

    @property
    def client(self) -> McpClientService:
        return self._client

    @property
    def registered_tools(self) -> set[str]:
        return set(self._registered)

    def _create_handler(
        self,
        server_id: str,
        tool_name: str,
    ):
        """Create a tool handler that delegates to the MCP client."""
        bridge = self  # capture for closure

        async def handler(inputs: ToolArgs, ctx: object) -> ToolOutput:
            # Prefer context-local client, fall back to bridge client
            client: McpClientService | None = None
            services = getattr(ctx, "services", None)
            if services is not None:
                client = getattr(services, "mcp_client", None)
            if client is None:
                client = bridge._client

            try:
                result = await client.call_tool(
                    server_id=server_id,
                    tool=tool_name,
                    arguments=dict(inputs) if inputs else {},
                )
                return ToolOutput(
                    success=True,
                    data={"result": result, "server_id": server_id, "tool": tool_name},
                )
            except Exception as e:
                logger.exception("MCP tool %s.%s failed", server_id, tool_name)
                return ToolOutput(
                    success=False,
                    error=f"MCP tool {server_id}.{tool_name} failed: {e}",
                )

        return handler

    def register_tool(
        self,
        *,
        server_id: str,
        tool_name: str,
        description: str,
        input_schema: dict[str, Any],
    ) -> str:
        """Register a single MCP tool in the global registry."""
        full_name = f"{server_id}.{tool_name}"
        if full_name in self._registered:
            return full_name

        registry = get_tool_registry()
        registry.register(
            name=full_name,
            handler=self._create_handler(server_id, tool_name),
            description=description or f"MCP tool: {full_name}",
            parameters=input_schema or {"type": "object", "properties": {}, "required": []},
            requires_privileged=False,
            mcp_server=server_id,
            mcp_tool=tool_name,
        )

        self._registered.add(full_name)
        logger.debug("Registered MCP tool: %s", full_name)
        return full_name


async def register_mcp_tools(
    *,
    mcp_client: McpClientService,
) -> list[str]:
    """Discover and register all tools from configured MCP servers.

    Args:
        mcp_client: The MCP client service to discover tools from.

    Returns
    -------
        List of registered tool names.
    """
    bridge = MCPToolBridge(mcp_client)
    registered: list[str] = []

    server_ids = mcp_client.list_server_ids()
    if not server_ids:
        return registered

    for server_id in server_ids:
        try:
            tools = await mcp_client.list_tools(server_id=server_id)
            for tool in tools:
                full_name = bridge.register_tool(
                    server_id=server_id,
                    tool_name=tool.name,
                    description=getattr(tool, "description", "") or "",
                    input_schema=getattr(tool, "inputSchema", {}) or {},
                )
                registered.append(full_name)
        except Exception as e:
            logger.warning(
                "Failed to list tools from MCP server '%s': %s",
                server_id, e, exc_info=True,
            )

    if registered:
        logger.info("Registered %d MCP tools from %d servers", len(registered), len(server_ids))

    return registered


__all__ = ["MCPToolBridge", "register_mcp_tools"]
