"""The ABOM MCP broker — enforcement moved OUT of the agent's process.

The runtime adapters (``abom.adapters``) are *cooperative*: they observe an agent
that chooses to route calls through them. A prompt-injected or buggy agent can
call a tool by another path and never trip the gate. The broker closes that gap
for MCP tools by relocating the gate to a **choke point the agent does not
control**: a proxy that sits between the agent (MCP client) and the upstream MCP
server. The agent gets its tools *from the broker*; every ``tools/call`` it makes
transits the gate before the broker forwards it upstream.

    agent ──MCP──▶  ABOM broker  ──MCP──▶  upstream MCP server
                     │  gate.check(tool)      (real tools live here)
                     │  DENY → JSON-RPC error, never forwarded
                     │  ALLOW → forward, notarize the decision
                     ▼
              persistent Merkle transparency log (signed tree heads)

What this is honest about (the scope caveat that makes it a *control*, not
theatre): the broker only mediates calls that flow through it. It makes
deny-by-default **structural for brokered MCP tools** — the agent cannot reach a
brokered tool without transiting the gate. It does NOT stop an agent that has a
Python interpreter and a network from calling out-of-band (raw HTTP, subprocess).
Closing *that* requires network-level egress containment, which is declared, not
provided here.

The core (``McpBroker.handle``) is pure over JSON-RPC dicts, so the whole
enforcement path is unit-testable with no sockets. ``forward`` is injected: an
in-memory callable in tests, an httpx POST to the upstream in production
(``connect_http``).
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable

from .gate import Action, Gate, Verdict

# JSON-RPC error code for a gate denial. -32000..-32099 is the reserved
# "server error" range in the JSON-RPC 2.0 spec; we use a stable code in it.
ABOM_DENIED_CODE = -32001

# MCP method names the broker treats specially.
M_TOOLS_CALL = "tools/call"
M_TOOLS_LIST = "tools/list"


# A forwarder takes an upstream-bound JSON-RPC request dict and returns the
# upstream's JSON-RPC response dict.
Forward = Callable[[dict], dict]


@dataclass
class BrokerStats:
    forwarded: int = 0
    denied: int = 0
    passthrough: int = 0  # non-tools/call methods forwarded untouched


@dataclass
class McpBroker:
    """Mediates MCP JSON-RPC between an agent and an upstream MCP server.

    ``gate`` decides ALLOW/DENY for each ``tools/call`` and notarizes it.
    ``forward`` sends a request upstream and returns its response.
    ``learn_tools`` (default False): when True, a ``tools/list`` response is used
    to note which tools the upstream actually offers (advisory — enforcement is
    always against the signed manifest via the gate, never the learned set).
    """

    gate: Gate
    forward: Forward
    learn_tools: bool = False
    stats: BrokerStats = field(default_factory=BrokerStats)
    upstream_tools: set[str] = field(default_factory=set)

    # -- helpers ---------------------------------------------------------------
    @staticmethod
    def _tool_name(request: dict) -> str | None:
        params = request.get("params") or {}
        return params.get("name")

    @staticmethod
    def _tool_args(request: dict) -> dict:
        params = request.get("params") or {}
        args = params.get("arguments")
        return args if isinstance(args, dict) else {}

    def _denied_response(self, request: dict, decision) -> dict:
        """A JSON-RPC error response the agent receives instead of the tool result."""
        return {
            "jsonrpc": "2.0",
            "id": request.get("id"),
            "error": {
                "code": ABOM_DENIED_CODE,
                "message": f"ABOM gate denied tool '{self._tool_name(request)}'",
                "data": {
                    "rule": decision.rule,
                    "reason": decision.reason,
                    "seq": decision.seq,
                    # the inclusion proof + signed head make the denial itself
                    # a notarized, verifiable event, not just a dropped call
                    "tree_head": decision.tree_head,
                    "head_signature": decision.head_signature,
                },
            },
        }

    # -- the core --------------------------------------------------------------
    def handle(self, request: dict) -> dict:
        """Handle one JSON-RPC request from the agent, returning the response the
        agent should receive.

        ``tools/call`` is gated before forwarding: a denied call is answered with
        a notarized JSON-RPC error and is NEVER sent upstream. Everything else is
        forwarded untouched (with ``tools/list`` optionally learned from).
        """
        method = request.get("method")

        if method == M_TOOLS_CALL:
            tool = self._tool_name(request)
            if not tool:
                return {
                    "jsonrpc": "2.0", "id": request.get("id"),
                    "error": {"code": -32602, "message": "tools/call missing params.name"},
                }
            decision = self.gate.check(Action(tool=tool, args=self._tool_args(request)))
            if decision.verdict is Verdict.DENY:
                self.stats.denied += 1
                return self._denied_response(request, decision)
            # allowed: forward upstream and return the real result
            self.stats.forwarded += 1
            return self.forward(request)

        # non-gated methods are forwarded verbatim
        response = self.forward(request)
        if method == M_TOOLS_LIST and self.learn_tools:
            self._learn(response)
        self.stats.passthrough += 1
        return response

    def _learn(self, response: dict) -> None:
        result = response.get("result") or {}
        for t in result.get("tools", []) or []:
            name = t.get("name") if isinstance(t, dict) else None
            if name:
                self.upstream_tools.add(name)


# --------------------------------------------------------------------------
# Production transport: forward over MCP Streamable HTTP (JSON-RPC 2.0 / POST).
# --------------------------------------------------------------------------
def connect_http(url: str, *, headers: dict | None = None, timeout: float = 30.0) -> Forward:
    """Build a ``forward`` that POSTs a JSON-RPC request to an upstream MCP URL
    and returns the parsed JSON-RPC response.

    Reuses the same Streamable-HTTP framing as ``abom.remote`` (application/json
    or a text/event-stream body). Requires httpx.
    """
    try:
        import httpx
    except ModuleNotFoundError as exc:  # pragma: no cover - environment guard
        raise RuntimeError(
            "the MCP broker needs httpx — install it with: pip install 'abom-cli[mcp]'"
        ) from exc

    from .remote import _parse_jsonrpc_body, _MCP_PROTOCOL_VERSION

    base_headers = {
        "Accept": "application/json, text/event-stream",
        "Content-Type": "application/json",
        "MCP-Protocol-Version": _MCP_PROTOCOL_VERSION,
    }
    if headers:
        base_headers.update(headers)

    client = httpx.Client(timeout=timeout, follow_redirects=True)

    def forward(request: dict) -> dict:
        resp = client.post(url, json=request, headers=base_headers)
        resp.raise_for_status()
        return _parse_jsonrpc_body(resp)

    return forward
