"""ABOM end-to-end demo — the whole product, A to Z, in one run.

Walks the complete pipeline a regulated team would use:

  1. SCAN     — build a signed Composition Manifest (what the agent IS).
  2. GATE     — enforce it: an undeclared tool is DENIED before it runs, and
                every decision is notarized in a durable Merkle transparency log.
  3. BROKER   — the same gate, relocated out of the agent's process: a denied
                tools/call never reaches the upstream MCP server.
  4. WITNESS  — an independent key cosigns the log and REFUSES a rewrite, so the
                record is tamper-evident against the operator, not just outsiders.
  5. ANCHOR   — the log root is timestamped by an RFC 3161 authority (mocked here
                so the demo runs offline) — a third-party existence proof.
  6. VERIFY   — the auditor's side: the cosignature checks out; a tampered one
                does not.

Everything is real code (the same functions the CLI calls); only the TSA network
call is stubbed so the demo runs with no dependencies and no network.

Run:  python3 demo/full_demo.py     (from the cli/ directory)
"""
from __future__ import annotations

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "src"))

from abom import (  # noqa: E402
    Action, Gate, McpBroker, Verdict, Witness,
    LocalSigner, anchor_head, bom,
)
from abom.merkle import PersistentMerkleLog  # noqa: E402
from abom import anchor as _anchor  # noqa: E402
from abom import sign as _sign  # noqa: E402

# ---- tiny presentation helpers (no deps) ------------------------------------
G, R, C, DIM, B, X = "\033[32m", "\033[31m", "\033[36m", "\033[90m", "\033[1m", "\033[0m"
TTY = sys.stdout.isatty()
def c(s, code): return f"{code}{s}{X}" if TTY else s
def step(n, title): print(f"\n{c(f'  {n}', C+B)}  {c(title, B)}")
def ok(s): print(f"     {c('✓', G)} {s}")
def no(s): print(f"     {c('✗', R)} {s}")
def note(s): print(f"       {c(s, DIM)}")


def main() -> int:
    workdir = os.path.join(HERE, "out")
    os.makedirs(workdir, exist_ok=True)
    logp = os.path.join(workdir, "full_demo.mlog")
    if os.path.exists(logp):
        os.remove(logp)  # fresh log each run
    for side in (".witness-state.json", ".cosig.json"):
        p = logp + side
        if os.path.exists(p):
            os.remove(p)

    print(c("\n  ABOM — end-to-end: scan → gate → broker → witness → anchor → verify", B))

    # 1. SCAN -----------------------------------------------------------------
    step("1", "SCAN — a signed Composition Manifest (what the agent is allowed)")
    manifest = bom.build_composition(
        agent={"name": "loan-doc-agent", "version": "1.4.0"},
        components=[
            {"type": "model", "name": "gpt-4o-mini"},
            {"type": "tool", "name": "lookup_customer"},
            {"type": "tool", "name": "read_kyc_doc"},
        ],
        controls={"egress": "deny-by-default"},
    )
    ok(f"signed manifest — {len(manifest['components'])} components, "
       f"ed25519 key {manifest['signature']['key_id']}")
    note(f"declares tools: lookup_customer, read_kyc_doc  (NOT wire_transfer)")

    # operator + witness hold DIFFERENT keys — that separation is the point
    operator = LocalSigner(_sign.load_or_create_key(os.path.join(workdir, "operator.pem")))
    witness_signer = LocalSigner(_sign.load_or_create_key(os.path.join(workdir, "witness.pem")))

    # 2. GATE -----------------------------------------------------------------
    step("2", "GATE — enforce the manifest, notarize every decision")
    gate = Gate(manifest, signer=operator, log_path=logp)
    d1 = gate.check(Action(tool="lookup_customer"))
    ok(f"lookup_customer  → {d1.verdict.value.upper()}  (declared)")
    d2 = gate.check(Action(tool="wire_transfer"))
    no(f"wire_transfer    → {d2.verdict.value.upper()}  ({d2.rule})")
    note(f"both decisions notarized · log root {gate.head()['root_hash'][:16]}…")

    # 3. BROKER ---------------------------------------------------------------
    step("3", "BROKER — the gate out of the agent's process (structural)")
    upstream_saw = []
    def fake_upstream(request):
        upstream_saw.append(request["params"]["name"])
        return {"jsonrpc": "2.0", "id": request["id"],
                "result": {"content": [{"type": "text", "text": "ran upstream"}]}}
    broker = McpBroker(gate=Gate(manifest, signer=operator, log_path=logp), forward=fake_upstream)
    for tool in ("read_kyc_doc", "wire_transfer"):
        resp = broker.handle({"jsonrpc": "2.0", "id": 1, "method": "tools/call",
                              "params": {"name": tool, "arguments": {}}})
        if "result" in resp:
            ok(f"tools/call {tool}  → forwarded to upstream")
        else:
            no(f"tools/call {tool}  → DENIED ({resp['error']['code']}) — upstream never saw it")
    note(f"upstream actually received: {upstream_saw}")

    # 4. WITNESS --------------------------------------------------------------
    step("4", "WITNESS — an independent key cosigns; a rewrite is refused")
    mlog = PersistentMerkleLog(logp)
    witness = Witness(witness_signer)
    cosig = witness.cosign(mlog.head().as_dict())
    ok(f"witness {witness.key_id} cosigned  size={cosig['tree_size']}, "
       f"root {cosig['root_hash'][:16]}…")
    note(f"witness key ≠ operator key ({witness.key_id} ≠ {operator.key_id()})")
    # operator tries to rewrite history to a smaller/forged log:
    from abom.witness import WitnessRejected
    forged = PersistentMerkleLog(os.path.join(workdir, "forged.mlog"))
    if forged.size == 0:
        forged.append(b"only-one-entry")
    try:
        witness.cosign(forged.head().as_dict())
        no("witness cosigned a rewrite — THIS SHOULD NOT HAPPEN")
    except WitnessRejected as e:
        ok(f"rewrite refused — {str(e).split(':')[0]}")

    # 5. ANCHOR ---------------------------------------------------------------
    step("5", "ANCHOR — RFC 3161 timestamp on the log root (third-party proof)")
    # stub the TSA network call so the demo runs offline; real usage hits a TSA URL
    _orig_post = _anchor._http_post
    _anchor._http_post = lambda url, tsq, timeout: b"DEMO-TSA-TOKEN"
    try:
        token = anchor_head(mlog.head().as_dict(), tsa_url="https://freetsa.org/tsr")
    finally:
        _anchor._http_post = _orig_post
    ok(f"root timestamped by {token['tsa_url']}")
    note(f"message imprint {token['message_imprint_sha256'][:16]}… · token stored for re-verification")
    note("(under eIDAS Art. 41 a qualified TSA gives this a legal presumption of time + integrity)")

    # 6. VERIFY ---------------------------------------------------------------
    step("6", "VERIFY — the auditor's side: no trust in the operator")
    good = Witness.verify_cosignature(cosig, witness_public_key_b64=witness.public_key_b64)
    (ok if good else no)(f"witness cosignature verifies → {good}")
    tampered = {**cosig, "root_hash": "00" * 32}
    bad = Witness.verify_cosignature(tampered, witness_public_key_b64=witness.public_key_b64)
    (ok if not bad else no)(f"tampered cosignature rejected → {not bad}")

    print(c("\n  ✓ End to end: composed, enforced, brokered, witnessed, anchored, verified.\n", G+B))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
