"""MCP adapter — connect an MCP server and register its tools in ToolRegistry.

This is the public entry-point for MCP client support.  ``connect_mcp()``
discovers tools from a running (or spawned) MCP server and registers each
one as a native ``ToolOutput``-returning handler, so the pathway VM sees
structured data instead of raw strings.

Usage::

    from infrastructure.mcp import connect_mcp

    conn = await connect_mcp(
        "servicem8",
        command=["python", "-m", "servicem8_mcp"],
    )
    # servicem8.list_records, servicem8.get_record, … are now in ToolRegistry
    await conn.close()
"""

from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING

from core.pathways.models.common import JSONSchema
from core.pathways.models.tool import ToolArgs, ToolOutput
from core.tools import ToolRegistry, get_tool_registry
from infrastructure.mcp.client import McpConnection

if TYPE_CHECKING:
    from mcp.types import CallToolResult

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Result conversion
# ---------------------------------------------------------------------------


def _to_tool_output(result: CallToolResult) -> ToolOutput:
    """Convert an MCP ``CallToolResult`` to a native ``ToolOutput``.

    Strategy:
    1. Extract all text content items from the result.
    2. If the server flagged an error, return ``ToolOutput(success=False)``.
    3. Try ``json.loads`` on the combined text — API-wrapping MCP servers
       almost always return JSON, so this succeeds in the common case and
       gives downstream nodes structured ``data`` for deterministic branching.
    4. Fall back to storing the raw string in ``data``.
    """
    texts = [c.text for c in result.content if hasattr(c, "text")]
    combined = "\n".join(texts)

    if result.isError:
        return ToolOutput(success=False, error=combined or "MCP tool returned an error")

    if not combined:
        return ToolOutput(success=True, data=None)

    # Try structured parse
    try:
        parsed = json.loads(combined)
        return ToolOutput(success=True, data=parsed)
    except (json.JSONDecodeError, ValueError):
        return ToolOutput(success=True, data=combined)


# ---------------------------------------------------------------------------
# Schema conversion
# ---------------------------------------------------------------------------


def _mcp_schema_to_json_schema(input_schema: dict) -> JSONSchema:
    """Convert an MCP tool's ``inputSchema`` to the albusOS ``JSONSchema`` model."""
    if not input_schema:
        return JSONSchema(properties={})
    return JSONSchema.from_dict(input_schema)


# ---------------------------------------------------------------------------
# Handler factory
# ---------------------------------------------------------------------------


def _make_handler(
    connection: McpConnection,
    mcp_tool_name: str,
):
    """Create a ``ToolHandler``-compatible async closure for one MCP tool.

    The returned function has the signature
    ``async (ToolArgs, ctx) -> ToolOutput`` expected by the ToolRegistry.
    """

    async def handler(args: ToolArgs, _ctx: object) -> ToolOutput:
        # Build the arguments dict from the Pydantic extras
        call_args: dict[str, object] = dict(args.__pydantic_extra__ or {})
        try:
            result = await connection.call_tool(mcp_tool_name, call_args)
            return _to_tool_output(result)
        except Exception as exc:
            logger.exception("MCP tool call failed: %s.%s", connection.namespace, mcp_tool_name)
            return ToolOutput(success=False, error=str(exc))

    return handler


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


async def connect_mcp(
    namespace: str,
    *,
    command: list[str] | None = None,
    url: str | None = None,
    env: dict[str, str] | None = None,
    registry: ToolRegistry | None = None,
) -> McpConnection:
    """Connect to an MCP server and register its tools in the ToolRegistry.

    Parameters
    ----------
    namespace
        Prefix for tool names.  An MCP tool named ``"list_records"`` becomes
        ``"{namespace}.list_records"`` in the registry.
    command
        Command + args to spawn a stdio-based MCP server
        (e.g. ``["python", "-m", "servicem8_mcp"]``).
    url
        URL of a running MCP server (SSE / Streamable HTTP).
    env
        Environment variables passed to the subprocess (stdio only).
    registry
        ToolRegistry to register into.  Defaults to the global registry.

    Returns
    -------
    McpConnection
        The live connection.  Call ``await conn.close()`` when done.

    Raises
    ------
    ImportError
        If the ``mcp`` package is not installed.
    ValueError
        If neither *command* nor *url* is provided.

    Examples
    --------
    ::

        conn = await connect_mcp("servicem8", command=["python", "-m", "servicem8_mcp"])
        # Tools are now registered: servicem8.list_records, servicem8.get_record, …
        await conn.close()
    """
    reg = registry if registry is not None else get_tool_registry()

    connection = McpConnection(
        namespace=namespace,
        command=command,
        url=url,
        env=env,
    )

    tools = await connection.connect()

    registered: list[str] = []
    for tool in tools:
        qualified_name = f"{namespace}.{tool.name}"
        schema = _mcp_schema_to_json_schema(tool.inputSchema)
        handler = _make_handler(connection, tool.name)

        reg.register(
            name=qualified_name,
            handler=handler,
            description=tool.description or f"MCP tool: {tool.name}",
            parameters=schema,
            category="mcp",
            skill_id=namespace,
        )
        registered.append(qualified_name)

    logger.info(
        "MCP [%s] registered %d tools: %s",
        namespace,
        len(registered),
        registered,
    )

    return connection


__all__ = ["connect_mcp"]
