"""Replay an agent run from ledger.jsonl, verify every token, print the DAG.

Model (spec-consistent)
-----------------------
Per run:

  * 1 ACT Phase 1 mandate (user → agent)
  * N ECT tokens — one per outgoing MCP HTTP request. Tool-call ECTs form
    a DAG via their ``pred`` field; session ECTs (initialize/tools-list)
    only point at the mandate.
  * 1 ACT Phase 2 record summarising the run (jti = mandate.jti per
    ACT §3.2).

The verifier re-runs the ietf-act and ietf-ect refimpls on each compact
form and prints both the coarse ACT summary and the fine-grained ECT DAG.

Run ``poc-verify [--ledger keys/ledger.jsonl] [--keys-dir keys]``.
"""

from __future__ import annotations

import argparse
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from act.crypto import ACTKeyResolver
from act.errors import ACTError
from act.verify import ACTVerifier

from ect.verify import verify as ect_verify, VerifyOptions

from .keys import build_ect_key_resolver, build_key_registry, load_identities


SERVER_IDENTITY_NAME = "mcp-server"


@dataclass
class Row:
    kind: str
    jti: str
    compact: str
    metadata: dict[str, Any]


def _read_ledger(path: Path) -> list[Row]:
    rows: list[Row] = []
    for line in path.read_text().splitlines():
        if not line.strip():
            continue
        obj = json.loads(line)
        rows.append(
            Row(
                kind=obj["kind"],
                jti=obj["jti"],
                compact=obj["compact"],
                metadata=obj.get("metadata", {}),
            )
        )
    return rows


def _fmt_jti(jti: str) -> str:
    return jti.split("-")[0]


def run(ledger_path: Path, keys_dir: Path) -> int:
    identities = load_identities(keys_dir)
    registry = build_key_registry(identities)
    resolver = ACTKeyResolver(registry=registry)
    ect_resolver = build_ect_key_resolver(identities)
    trusted_issuers = {ident.name for ident in identities.values()}

    verifier = ACTVerifier(
        resolver,
        verifier_id=SERVER_IDENTITY_NAME,
        trusted_issuers=trusted_issuers,
    )

    rows = _read_ledger(ledger_path)
    mandates = [r for r in rows if r.kind == "mandate"]
    records = [r for r in rows if r.kind == "record"]
    ect_rows = [r for r in rows if r.kind == "ect"]
    if len(mandates) != 1:
        raise SystemExit(
            f"expected exactly one mandate, got {len(mandates)} in {ledger_path}"
        )
    if len(records) != 1:
        raise SystemExit(
            f"expected exactly one record, got {len(records)} in {ledger_path}"
        )

    try:
        mandate = verifier.verify_mandate(mandates[0].compact, check_sub=False)
    except ACTError as e:
        raise SystemExit(f"mandate verification failed: {e}")
    print(f"mandate  verified   jti={_fmt_jti(mandate.jti)}")

    # ECT verification — includes refimpl DAG walk when we supply a store.
    # We don't supply one here because ECTStore would need cross-run scoping.
    # Each ECT still passes its own Section-7 verification individually.
    ect_parsed: list[Any] = []
    ect_sessions = 0
    ect_tool_calls = 0
    for row in ect_rows:
        parsed = ect_verify(
            row.compact,
            VerifyOptions(verifier_id=SERVER_IDENTITY_NAME, resolve_key=ect_resolver),
        )
        ect_parsed.append(parsed)
        if row.metadata.get("session_only"):
            ect_sessions += 1
        else:
            ect_tool_calls += 1
    print(
        f"ects     verified   n={len(ect_parsed)} "
        f"(tool-calls={ect_tool_calls}, session={ect_sessions})"
    )

    # Final ACT record — verify without the DAG store (pred=[] for our model).
    try:
        record = verifier.verify_record(records[0].compact, store=None)
    except ACTError as e:
        raise SystemExit(f"record verification failed: {e}")
    if record.jti != mandate.jti:
        raise SystemExit(
            f"record.jti {record.jti!r} != mandate.jti {mandate.jti!r} "
            "— ACT §3.2 violation"
        )
    print(f"record   verified   jti={_fmt_jti(record.jti)} status={record.status}")

    # Cross-check: ECT DAG well-formedness within this run.
    known_jtis = {mandate.jti} | {p.payload.jti for p in ect_parsed}
    dangling = 0
    for p in ect_parsed:
        for pred in p.payload.pred:
            if pred not in known_jtis:
                dangling += 1
    if dangling:
        raise SystemExit(f"ECT DAG has {dangling} dangling predecessor ref(s)")
    print("ect-dag  wellformed  every pred is the mandate or a prior ECT")

    # ---- Render ------------------------------------------------------------
    print()
    print("Run")
    print("===")
    print(f"  mandate {_fmt_jti(mandate.jti)}  task={mandate.task.purpose!r}")
    print(f"    iss={mandate.iss}  sub={mandate.sub}  aud={mandate.aud}")
    print(f"    cap={[c.action for c in mandate.cap]}")
    print()
    print("Tool-call ECT DAG:")
    tool_only = [
        p
        for p, row in zip(ect_parsed, ect_rows)
        if not row.metadata.get("session_only")
    ]
    if not tool_only:
        print("  (none — model called no tools)")
    for p in tool_only:
        preds = [_fmt_jti(x) for x in p.payload.pred]
        print(
            f"  ect {_fmt_jti(p.payload.jti)}  exec_act={p.payload.exec_act}  "
            f"pred={preds}"
        )
    print()
    print("ACT Phase 2 record:")
    print(f"  jti={_fmt_jti(record.jti)}  exec_act={record.exec_act}")
    print(f"  status={record.status}  pred={list(record.pred)}")
    print(f"  inp_hash={record.inp_hash}")
    print(f"  out_hash={record.out_hash}")
    return 0


def main() -> None:
    parser = argparse.ArgumentParser(description="Verify ACT+ECT PoC ledger")
    parser.add_argument(
        "--ledger", default=os.environ.get("POC_LEDGER", "keys/ledger.jsonl")
    )
    parser.add_argument("--keys-dir", default=os.environ.get("POC_KEYS_DIR", "keys"))
    args = parser.parse_args()
    raise SystemExit(run(Path(args.ledger), Path(args.keys_dir)))


if __name__ == "__main__":
    main()
