"""Tests for the ABOM MCP broker — the pivot's core property.

The thesis: a denied tools/call is answered with a notarized JSON-RPC error and
is NEVER forwarded upstream. Enforcement is structural for brokered tools — the
agent cannot reach a brokered tool without transiting the gate. These tests prove
that with an in-memory forwarder (no sockets), so the enforcement path is exact.
"""
import pytest

from abom import sign
from abom.bom import build_composition
from abom.broker import ABOM_DENIED_CODE, McpBroker, M_TOOLS_CALL, M_TOOLS_LIST
from abom.gate import Gate
from abom.signer import LocalSigner


@pytest.fixture
def manifest(tmp_path):
    key = sign.load_or_create_key(tmp_path / "owner.pem")
    return build_composition(
        agent={"name": "brokered-agent", "version": "1.0.0"},
        components=[{"type": "tool", "name": "search_orders"},
                    {"type": "tool", "name": "read_docs"}],
        controls={},
        key=key,
    )


@pytest.fixture
def signer(tmp_path):
    return LocalSigner(sign.load_or_create_key(tmp_path / "owner.pem"))


def _spy_forward():
    """A forwarder that records what reached upstream and returns a canned result."""
    seen = []

    def forward(request):
        seen.append(request)
        return {"jsonrpc": "2.0", "id": request.get("id"),
                "result": {"content": [{"type": "text", "text": "upstream ok"}]}}

    forward.seen = seen
    return forward


def _call(tool, *, id=1, args=None):
    return {"jsonrpc": "2.0", "id": id, "method": M_TOOLS_CALL,
            "params": {"name": tool, "arguments": args or {}}}


# --------------------------------------------------------------------------
# the structural property
# --------------------------------------------------------------------------
def test_declared_tool_is_forwarded(manifest, signer):
    fwd = _spy_forward()
    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=fwd)

    resp = broker.handle(_call("search_orders"))

    assert "result" in resp
    assert resp["result"]["content"][0]["text"] == "upstream ok"
    assert len(fwd.seen) == 1  # it reached upstream
    assert broker.stats.forwarded == 1


def test_undeclared_tool_is_denied_and_never_forwarded(manifest, signer):
    """The whole point: a denied call does not transit to the upstream server."""
    fwd = _spy_forward()
    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=fwd)

    resp = broker.handle(_call("wire_transfer", id=7))

    assert "error" in resp
    assert resp["error"]["code"] == ABOM_DENIED_CODE
    assert resp["id"] == 7
    assert fwd.seen == []  # NEVER forwarded — structural enforcement
    assert broker.stats.denied == 1
    assert broker.stats.forwarded == 0


def test_denial_is_notarized_with_verifiable_inclusion_proof(manifest, signer):
    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=_spy_forward())
    resp = broker.handle(_call("wire_transfer"))

    data = resp["error"]["data"]
    assert data["rule"] == "tool_not_in_manifest"
    assert "tree_head" in data and "head_signature" in data
    # the decision is a real leaf in the gate's transparency log
    assert broker.gate.verify()["valid"]


def test_denial_persisted_to_durable_log(manifest, signer, tmp_path):
    logp = tmp_path / "broker.mlog"
    gate = Gate(manifest, signer=signer, log_path=logp)
    broker = McpBroker(gate=gate, forward=_spy_forward())

    broker.handle(_call("search_orders"))   # allow
    broker.handle(_call("wire_transfer"))   # deny

    # both decisions are on disk and survive a "restart"
    from abom.merkle import PersistentMerkleLog
    reloaded = PersistentMerkleLog(logp)
    assert reloaded.size == 2


# --------------------------------------------------------------------------
# passthrough + learning
# --------------------------------------------------------------------------
def test_non_tool_methods_pass_through(manifest, signer):
    fwd = _spy_forward()
    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=fwd)

    req = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}
    broker.handle(req)
    assert fwd.seen == [req]
    assert broker.stats.passthrough == 1


def test_tools_list_learning_records_upstream_tools(manifest, signer):
    def forward(request):
        return {"jsonrpc": "2.0", "id": request.get("id"),
                "result": {"tools": [{"name": "search_orders"}, {"name": "extra_tool"}]}}

    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=forward, learn_tools=True)
    broker.handle({"jsonrpc": "2.0", "id": 1, "method": M_TOOLS_LIST, "params": {}})
    assert broker.upstream_tools == {"search_orders", "extra_tool"}


def test_tools_call_missing_name_is_a_protocol_error(manifest, signer):
    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=_spy_forward())
    resp = broker.handle({"jsonrpc": "2.0", "id": 1, "method": M_TOOLS_CALL, "params": {}})
    assert resp["error"]["code"] == -32602


def test_broker_refuses_untrusted_manifest_via_gate(manifest, signer, tmp_path):
    """The broker inherits the gate's trust-pinning: an attacker-re-signed
    manifest cannot even build the gate the broker runs on."""
    owner_kid = manifest["signature"]["key_id"]
    attacker = sign.load_or_create_key(tmp_path / "attacker.pem")
    manifest["components"].append({"type": "tool", "name": "wire_transfer"})
    forged = sign.sign_obj({k: v for k, v in manifest.items() if k != "signature"}, attacker)

    with pytest.raises(ValueError, match="untrusted"):
        Gate(forged, trusted_keys={owner_kid})


# --------------------------------------------------------------------------
# connect_http — the production forwarder (mocked transport, no network)
# --------------------------------------------------------------------------
def test_connect_http_forwards_and_parses(monkeypatch, manifest, signer):
    httpx = pytest.importorskip("httpx")
    from abom import broker as broker_mod

    def handler(request):
        import json as _json
        req = _json.loads(request.content)
        return httpx.Response(200, json={
            "jsonrpc": "2.0", "id": req.get("id"),
            "result": {"content": [{"type": "text", "text": "real upstream"}]},
        })

    transport = httpx.MockTransport(handler)
    real_client = httpx.Client
    monkeypatch.setattr(httpx, "Client",
                        lambda *a, **k: real_client(*a, **{**k, "transport": transport}))

    forward = broker_mod.connect_http("https://upstream.example.com/mcp")
    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=forward)

    resp = broker.handle(_call("search_orders"))
    assert resp["result"]["content"][0]["text"] == "real upstream"


def test_connect_http_denied_call_never_hits_the_network(monkeypatch, manifest, signer):
    httpx = pytest.importorskip("httpx")
    from abom import broker as broker_mod

    hits = []

    def handler(request):
        hits.append(request.url)
        return httpx.Response(200, json={"jsonrpc": "2.0", "id": 1, "result": {}})

    transport = httpx.MockTransport(handler)
    real_client = httpx.Client
    monkeypatch.setattr(httpx, "Client",
                        lambda *a, **k: real_client(*a, **{**k, "transport": transport}))

    forward = broker_mod.connect_http("https://upstream.example.com/mcp")
    broker = McpBroker(gate=Gate(manifest, signer=signer), forward=forward)

    resp = broker.handle(_call("wire_transfer"))
    assert resp["error"]["code"] == ABOM_DENIED_CODE
    assert hits == []  # the denied call never reached the upstream server
