#!/usr/bin/env python3
"""Validate the focuschain-ab treatment smoke run.

Pass criteria:
- every step's tool_use input carries a non-empty task_notes field;
- zero context edits (no image_trim/compaction context_events);
- zero cache invalidation events (cache_read collapse) beyond step 1;
- cache hit rate comparable to keep_all (>= 0.80).
"""
import json
import sys
from pathlib import Path

BASE = Path("/srv/controller-data/results/focuschain-ab-smoke")

ok = True
tasks = sorted(BASE.glob("*/claude_computer_use/screenshot/*/*/*/traj.jsonl"))
if not tasks:
    print("NO TRAJECTORIES FOUND"); sys.exit(1)

for traj in tasks:
    task_id = traj.parent.name
    steps = {}
    notes_missing = 0
    rows = 0
    for line in traj.read_text(errors="replace").splitlines():
        try:
            row = json.loads(line)
        except Exception:
            continue
        rows += 1
        action = row.get("action") or {}
        inp = action.get("input") or {}
        if isinstance(inp, dict) and not (inp.get("task_notes") or "").strip():
            notes_missing += 1
        u = action.get("usage") or {}
        sn = row.get("step_num")
        if u and sn is not None and sn not in steps:
            steps[sn] = u

    per = []
    for sn in sorted(steps):
        u = steps[sn]
        inp = int(u.get("input_tokens") or 0)
        cr = int(u.get("cache_creation_input_tokens") or 0)
        rd = int(u.get("cache_read_input_tokens") or 0)
        per.append({"step": sn, "prompt": inp + cr + rd, "read": rd,
                    "event": u.get("context_event") or ""})

    edits = [p for p in per if p["event"]]
    inval = sum(1 for a, b in zip(per, per[1:])
                if a["prompt"] > 2000 and b["read"] < 0.5 * a["prompt"])
    tail = [p for p in per if p["step"] >= 2]
    denom = sum(p["prompt"] for p in tail)
    hit = sum(p["read"] for p in tail) / denom if denom else 0.0
    result_file = traj.parent / "result.txt"
    score = result_file.read_text().strip() if result_file.exists() else "n/a"

    status = "PASS"
    if edits or inval or hit < 0.80 or notes_missing:
        status = "FAIL"
        ok = False
    print(f"{status} {task_id}: steps={len(per)} score={score} edits={len(edits)} "
          f"invalidations={inval} hit_rate={hit:.3f} rows_missing_notes={notes_missing}/{rows}")
    for e in edits:
        print(f"    edit event at step {e['step']}: {e['event']}")

sys.exit(0 if ok else 1)
