"""MCP (Model Context Protocol) server for SwarmNode.

Exposes a SwarmNode's capabilities as MCP tools over an HTTP/JSON-RPC 2.0
endpoint so that any MCP-compatible client (e.g. Claude Desktop, Cursor) can
discover and invoke the node's capabilities directly.

Endpoints:
    POST /mcp   — JSON-RPC 2.0 handler supporting ``tools/list`` and
                  ``tools/call`` methods.
    GET  /mcp/health — Health check; returns ``{"status": "ok", "toolCount": N}``.

Usage::

    from meshpocalypse_semantic import SwarmNode, SwarmNodeConfig
    import asyncio

    config = SwarmNodeConfig(agent_id="my-agent", capabilities=["summarise", "search"])
    node = SwarmNode(config)
    await node.connect()
    await node.as_mcp_server(port=3001)  # blocks until cancelled
"""

from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass, field
from typing import Any

try:
    from aiohttp import web
except ImportError as exc:  # pragma: no cover
    raise ImportError(
        "aiohttp is required for MCP server support. "
        "Install it with: pip install aiohttp"
    ) from exc

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

_JSONRPC_VERSION = "2.0"

_MCP_TOOL_INPUT_SCHEMA: dict[str, Any] = {
    "type": "object",
    "properties": {
        "input": {
            "type": "string",
            "description": "Input payload for the capability (JSON-encoded string).",
        }
    },
    "required": ["input"],
}


@dataclass
class MCPTool:
    """A single MCP tool derived from a SwarmNode capability.

    Attributes:
        name: Tool name (matches the capability string).
        description: Human-readable description, either from
            ``tool_descriptions`` or a generated default.
        input_schema: JSON Schema describing the tool's input parameters.
    """

    name: str
    description: str
    input_schema: dict[str, Any] = field(default_factory=lambda: dict(_MCP_TOOL_INPUT_SCHEMA))

    def to_dict(self) -> dict[str, Any]:
        """Serialise to MCP tool-list format.

        Returns:
            Dict with ``name``, ``description``, and ``inputSchema`` keys.
        """
        return {
            "name": self.name,
            "description": self.description,
            "inputSchema": self.input_schema,
        }


@dataclass
class MCPServerConfig:
    """Configuration for the MCP HTTP server.

    Attributes:
        agent_id: Agent ID of the SwarmNode (used in error messages).
        capabilities: Capability names advertised by the node.
        nats_client: Live NATS client for routing ``tools/call`` requests.
            May be ``None`` if the node is not connected; calls will return
            an error response in that case.
        tool_descriptions: Optional mapping from capability name to a custom
            description string. Capabilities not present in this mapping get
            an auto-generated description.
        host: Bind host for the HTTP server. Defaults to ``"0.0.0.0"``.
        port: TCP port for the HTTP server. Defaults to ``3001``.
    """

    agent_id: str
    capabilities: list[str]
    nats_client: Any  # nats.aio.client.Client | None
    tool_descriptions: dict[str, str] = field(default_factory=dict)
    host: str = "0.0.0.0"
    port: int = 3001


# ---------------------------------------------------------------------------
# JSON-RPC helpers
# ---------------------------------------------------------------------------


def _jsonrpc_result(request_id: Any, result: Any) -> dict[str, Any]:
    """Build a JSON-RPC 2.0 success response.

    Args:
        request_id: The ``id`` from the incoming request.
        result: The result payload.

    Returns:
        JSON-RPC 2.0 response dict.
    """
    return {"jsonrpc": _JSONRPC_VERSION, "id": request_id, "result": result}


def _jsonrpc_error(
    request_id: Any,
    code: int,
    message: str,
    data: Any = None,
) -> dict[str, Any]:
    """Build a JSON-RPC 2.0 error response.

    Args:
        request_id: The ``id`` from the incoming request (may be ``None``).
        code: JSON-RPC error code.
        message: Short error message.
        data: Optional additional error data.

    Returns:
        JSON-RPC 2.0 error response dict.
    """
    error: dict[str, Any] = {"code": code, "message": message}
    if data is not None:
        error["data"] = data
    return {"jsonrpc": _JSONRPC_VERSION, "id": request_id, "error": error}


# ---------------------------------------------------------------------------
# Tool generation
# ---------------------------------------------------------------------------


def _build_tools(config: MCPServerConfig) -> list[MCPTool]:
    """Generate an MCPTool for every declared capability.

    Args:
        config: MCP server configuration.

    Returns:
        Ordered list of :class:`MCPTool` instances.
    """
    tools: list[MCPTool] = []
    for cap in config.capabilities:
        description = config.tool_descriptions.get(
            cap,
            f"Invoke the '{cap}' capability on agent '{config.agent_id}'.",
        )
        tools.append(MCPTool(name=cap, description=description))
    return tools


# ---------------------------------------------------------------------------
# Request handlers
# ---------------------------------------------------------------------------


async def _handle_tools_list(
    tools: list[MCPTool],
    request_id: Any,
) -> dict[str, Any]:
    """Handle the ``tools/list`` MCP method.

    Args:
        tools: Pre-generated tool list.
        request_id: JSON-RPC request id.

    Returns:
        JSON-RPC result containing ``{"tools": [...]}``.
    """
    return _jsonrpc_result(request_id, {"tools": [t.to_dict() for t in tools]})


async def _handle_tools_call(
    params: dict[str, Any],
    tools: list[MCPTool],
    config: MCPServerConfig,
    request_id: Any,
) -> dict[str, Any]:
    """Handle the ``tools/call`` MCP method.

    Routes the call through the NATS connection to the agent's capability
    subject (``mesh.capability.<name>``). Returns an error if the node is
    not connected or the capability is unknown.

    Args:
        params: JSON-RPC params dict (must contain ``name`` and optionally
            ``arguments``).
        tools: Pre-generated tool list for capability validation.
        config: MCP server configuration (for NATS client access).
        request_id: JSON-RPC request id.

    Returns:
        JSON-RPC success or error response dict.
    """
    tool_name: str | None = params.get("name")
    if not tool_name:
        return _jsonrpc_error(request_id, -32602, "Missing 'name' in params")

    known_names = {t.name for t in tools}
    if tool_name not in known_names:
        return _jsonrpc_error(
            request_id,
            -32602,
            f"Unknown tool '{tool_name}'",
            data={"availableTools": sorted(known_names)},
        )

    if config.nats_client is None:
        return _jsonrpc_error(
            request_id,
            -32000,
            "SwarmNode is not connected to NATS — call await node.connect() first",
        )

    arguments: dict[str, Any] = params.get("arguments", {})
    subject = f"mesh.capability.{tool_name}"
    payload_bytes = json.dumps(
        {
            "agentId": config.agent_id,
            "tool": tool_name,
            "arguments": arguments,
        }
    ).encode("utf-8")

    try:
        response = await config.nats_client.request(subject, payload_bytes, timeout=10.0)
        result_data: Any = json.loads(response.data.decode("utf-8"))
        return _jsonrpc_result(request_id, {"content": result_data})
    except json.JSONDecodeError as exc:
        return _jsonrpc_error(request_id, -32603, f"Invalid JSON from NATS: {exc}")
    except asyncio.TimeoutError:
        return _jsonrpc_error(
            request_id,
            -32000,
            f"NATS request timed out for capability '{tool_name}'",
        )
    except Exception as exc:  # noqa: BLE001 — catch-all after specific ones
        logger.exception("Unexpected error routing tools/call for '%s'", tool_name)
        return _jsonrpc_error(request_id, -32603, f"Internal error: {exc}")


# ---------------------------------------------------------------------------
# aiohttp application factory
# ---------------------------------------------------------------------------


def _build_app(config: MCPServerConfig) -> web.Application:
    """Build and configure the aiohttp web application.

    Registers ``POST /mcp`` and ``GET /mcp/health`` routes.

    Args:
        config: MCP server configuration.

    Returns:
        Configured :class:`aiohttp.web.Application`.
    """
    tools = _build_tools(config)
    app = web.Application()

    async def post_mcp(request: web.Request) -> web.Response:
        """Handle POST /mcp — JSON-RPC 2.0 dispatcher.

        Args:
            request: Incoming HTTP request.

        Returns:
            JSON HTTP response.
        """
        try:
            body: dict[str, Any] = await request.json()
        except (json.JSONDecodeError, ValueError) as exc:
            return web.Response(
                content_type="application/json",
                status=400,
                text=json.dumps(
                    _jsonrpc_error(None, -32700, f"Parse error: {exc}")
                ),
            )

        request_id = body.get("id")
        method: str = body.get("method", "")
        params: dict[str, Any] = body.get("params") or {}

        if method == "tools/list":
            result = await _handle_tools_list(tools, request_id)
        elif method == "tools/call":
            result = await _handle_tools_call(params, tools, config, request_id)
        else:
            result = _jsonrpc_error(
                request_id, -32601, f"Method not found: '{method}'"
            )

        return web.Response(
            content_type="application/json",
            text=json.dumps(result),
        )

    async def get_health(request: web.Request) -> web.Response:  # noqa: ARG001
        """Handle GET /mcp/health.

        Args:
            request: Incoming HTTP request (unused).

        Returns:
            JSON HTTP response with status and tool count.
        """
        return web.Response(
            content_type="application/json",
            text=json.dumps({"status": "ok", "toolCount": len(tools)}),
        )

    app.router.add_post("/mcp", post_mcp)
    app.router.add_get("/mcp/health", get_health)
    return app


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


async def start_mcp_server(config: MCPServerConfig) -> None:
    """Start the MCP HTTP server and run until the current task is cancelled.

    Binds to ``config.host:config.port``, serves requests, and performs a
    clean shutdown (draining in-flight requests) when an
    :class:`asyncio.CancelledError` is received.

    Args:
        config: Server configuration including host, port, capabilities, and
            the optional NATS client for routing.

    Raises:
        OSError: If the port is already in use or the bind address is invalid.
        asyncio.CancelledError: Re-raised after clean shutdown so callers can
            chain cancellation correctly.
    """
    app = _build_app(config)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, config.host, config.port)
    await site.start()

    logger.info(
        "[MCPServer] Listening on http://%s:%d/mcp (agent=%s, tools=%d)",
        config.host,
        config.port,
        config.agent_id,
        len(config.capabilities),
    )

    try:
        # Park here until cancelled
        await asyncio.get_event_loop().create_future()
    except asyncio.CancelledError:
        logger.info("[MCPServer] Shutting down — draining connections")
        await runner.cleanup()
        raise
