#!/usr/bin/env python3
"""Paired speed benchmark: PortMem review vs an LLM agent verifying the same memo.

Backs (or refutes) the site's Speed tile ("a full report verified in a fraction of
the time an LLM agent takes"). Same judge model on every side (the prod-parity
routing chain, Haiku 4.5), same document, same memo, faithfulness only (no B pass)
so the comparison is claim-verification wall-clock, nothing else.

Sides:
  portmem   — PortMem.review() (retrieval-scoped evidence, concurrent claim judging).
  agent_1pass — the FASTEST possible agent: one long-context call, full document +
                memo, verify every claim in one pass. Lower bound for any agent.
  agent_loop  — the typical agent pattern: extract claims (1 call), then verify each
                claim sequentially against the full document (n calls). What a
                naive "just use an agent loop" build actually does.

Run (local stack, portmem-engine-pg on :5434 + routing.haiku.yaml):
  cd product1 && python3 ../eval/speed_benchmark.py
"""
from __future__ import annotations

import asyncio
import json
import re
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "product1"))

from dogfood_domain import ENV_DEFAULTS, _llm, _pm, score_seeded  # noqa: E402


def score_agent(raw: str, errors: list[dict]) -> dict:
    """Same corrected matching rule as score_seeded (diff-token, word-boundary,
    vocabulary overlap — see dogfood_domain._claim_matches; the old any-number
    fragment rule degenerated to single digits, 2026-08-08 audit), applied to an
    agent transcript line by line."""
    from dogfood_domain import _claim_matches
    lines = [ln.strip() for ln in raw.splitlines() if ln.strip()]
    caught = fs = um = 0
    for e in errors:
        hits = [ln for ln in lines if _claim_matches(ln, e)]
        if not hits:
            um += 1
        elif any(re.search(r"CONTRADICTED|UNCERTAIN", ln, re.I) for ln in hits):
            caught += 1
        else:
            fs += 1
    return {"caught": caught, "false_supported": fs, "unmatched": um, "n": len(errors)}

DOC = ROOT / "eval" / "dogfood" / "corpora" / "eu_ai_act_ch1_3.txt"
MEMO = ROOT / "eval" / "dogfood" / "EU_AI_Act_Ch_I-III" / "report_corrupted.txt"
ERRORS = ROOT / "eval" / "dogfood" / "EU_AI_Act_Ch_I-III" / "errors.json"
SOURCE_ID = "EU AI Act Ch I-III"
MAX_CLAIMS = 25  # override with --max-claims
OUT = ROOT / "eval" / "dogfood" / "speed_benchmark_results.json"  # override with --out

EXTRACT_PROMPT = """List the distinct, checkable factual claims in the memo below, one per line,
at most {n}. No commentary.

MEMO:
{memo}"""

VERIFY_ONE_PROMPT = """You are verifying a claim from an analyst memo against the governing document.
Answer with exactly one word first — SUPPORTED, CONTRADICTED, or UNCERTAIN — then one
sentence citing where in the document you looked.

DOCUMENT:
{doc}

CLAIM: {claim}"""

VERIFY_ALL_PROMPT = """You are verifying an analyst memo against the governing document. For EVERY
distinct factual claim in the memo (up to {n}), output one line:
<verdict SUPPORTED|CONTRADICTED|UNCERTAIN> | <the claim> | <one-clause justification>

DOCUMENT:
{doc}

MEMO:
{memo}"""


async def main() -> None:
    import argparse
    import os
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-claims", type=int, default=MAX_CLAIMS)
    ap.add_argument("--out", default=str(OUT))
    a = ap.parse_args()
    max_claims = a.max_claims
    out_path = Path(a.out)
    tag = out_path.stem.replace("speed_benchmark_results", "") or "_base"
    for k, v in ENV_DEFAULTS.items():
        os.environ.setdefault(k, v)
    from run_local_verify import _load_env
    _load_env()

    doc = DOC.read_text()
    memo = MEMO.read_text()
    pm = await _pm()
    results: dict[str, dict] = {}

    # ── side 1: PortMem review (A only) ──────────────────────────────────────
    t0 = time.monotonic()
    rv = await pm.review(memo, max_claims=max_claims, max_results=8,
                         source=SOURCE_ID, check_completeness=False)
    dt = time.monotonic() - t0
    f = rv["faithfulness"]
    seeds = json.loads(ERRORS.read_text())
    sc = score_seeded(rv, seeds)
    results["portmem"] = {
        "wall_s": round(dt, 1), "n_claims": f.get("n_claims"),
        "verdicts": {k: f.get(f"n_{k}") for k in
                     ("supported", "contradicted", "uncertain", "not_found")},
        "seeded": {"caught": sum(x["caught"] for x in sc),
                   "false_supported": sum(x["false_supported"] for x in sc),
                   "unmatched": sum(x["unmatched"] for x in sc), "n": len(sc)},
        "seeded_detail": sc,
        "claims": [{"claim": c["claim"], "verdict": c["verdict"]}
                   for c in f.get("claims", [])],
    }
    print(f"portmem: {dt:.1f}s over {f.get('n_claims')} claims", flush=True)

    # ── side 2: single long-context pass (agent lower bound) ─────────────────
    t0 = time.monotonic()
    out = await _llm(pm, VERIFY_ALL_PROMPT.format(n=max_claims, doc=doc, memo=memo),
                     max_tokens=8192)
    dt = time.monotonic() - t0
    (ROOT / "eval" / "dogfood" / f"speed_agent_1pass_raw{tag}.txt").write_text(out)
    lines = [ln for ln in out.splitlines() if re.search(
        r"\b(SUPPORTED|CONTRADICTED|UNCERTAIN)\b", ln, re.I)]
    results["agent_1pass"] = {"wall_s": round(dt, 1), "n_claims": len(lines),
                              "seeded": score_agent(out, seeds)}
    print(f"agent_1pass: {dt:.1f}s over {len(lines)} claims "
          f"seeded={results['agent_1pass']['seeded']}", flush=True)

    # ── side 3: sequential claim-by-claim loop (typical agent) ───────────────
    t0 = time.monotonic()
    ext = await _llm(pm, EXTRACT_PROMPT.format(n=max_claims, memo=memo), max_tokens=2048)
    claims = [c.strip("-• \t") for c in ext.splitlines() if len(c.strip()) > 15][:max_claims]
    verdicts = []
    for c in claims:
        v = await _llm(pm, VERIFY_ONE_PROMPT.format(doc=doc, claim=c), max_tokens=300)
        verdicts.append(v.split()[0].upper() if v.split() else "?")
    dt = time.monotonic() - t0
    transcript = "\n".join(f"{v} | {c}" for v, c in zip(verdicts, claims))
    (ROOT / "eval" / "dogfood" / f"speed_agent_loop_raw{tag}.txt").write_text(transcript)
    results["agent_loop"] = {"wall_s": round(dt, 1), "n_claims": len(claims),
                             "verdict_mix": {k: verdicts.count(k) for k in set(verdicts)},
                             "seeded": score_agent(transcript, seeds)}
    print(f"agent_loop: {dt:.1f}s over {len(claims)} claims "
          f"seeded={results['agent_loop']['seeded']}", flush=True)

    results["_meta"] = {
        "doc_chars": len(doc), "memo_chars": len(memo), "max_claims": max_claims,
        "judge": "prod-parity routing chain (routing.haiku.yaml -> Haiku 4.5)",
        "note": ("faithfulness only on every side, on the CORRUPTED memo so speed and "
                 "catch-rate are measured together; PortMem time includes claim "
                 "extraction + retrieval + concurrent judging + calibration; the agent "
                 "sides get the full document handed to them free (no search cost), "
                 "which is GENEROUS to the agent — a real agent must find its context"),
    }
    out_path.write_text(json.dumps(results, indent=1, default=str))
    p, a1, a2 = (results[k]["wall_s"] for k in ("portmem", "agent_1pass", "agent_loop"))
    print(f"\nportmem {p}s | agent_1pass {a1}s ({a1/p:.1f}x) | agent_loop {a2}s ({a2/p:.1f}x)")


if __name__ == "__main__":
    asyncio.run(main())