#!/usr/bin/env python3
"""Build the shared 40-task OSWorld pack for the focuschain-ab experiment.

Selection rules (deterministic, no RNG):
- Universe: evaluation_examples/test_nogdrive.json (361 tasks).
- Step counts: Pointer's published sonnet-4-6 run, max screenshot step per task
  from trials/trace-forensics/data/pointer_png_index.csv.
- Strata: medium = 9-15 steps, long = >=16 steps; 20 tasks each.
- Exclusions:
  * every task flagged in trials/trace-forensics/findings/benchmark-integrity.md
    with a non-AGENT_FAIL category (broken evaluators, live-web drift,
    infeasible-missed, gate false positives) — from both_fail_classification.csv;
  * every infeasible-type task (evaluator func contains "infeasible");
  * belt-and-braces: the explicit ID list from the findings doc.
- Within each stratum: round-robin across domains (domains sorted
  alphabetically, tasks within a domain sorted by descending step count then
  task_id) so the pack is domain-diverse and skewed toward the longer end of
  each stratum.

Outputs (to trials/shared/osworld-40-pack/):
- pack.json        OSWorld --test_all_meta_path format: {domain: [task_ids]}
- pack_detail.csv  domain,task_id,pointer_steps,stratum
- pack.md          human-readable documentation
"""
import csv
import json
from collections import defaultdict
from pathlib import Path

ROOT = Path("/srv/controller-data")
OSW = ROOT / "repos/OSWorld"
OUT = ROOT / "trials/shared/osworld-40-pack"

# Explicit flagged IDs (8-char prefixes) from benchmark-integrity.md A1 table
# and prose, kept as a belt-and-braces overlay on the classification CSV.
FLAGGED_PREFIXES = {
    # broken evaluator CONFIRMED
    "b4f95342", "42f4d1c7", "f3b19d1e", "9f935cce", "08c73485", "5990457f",
    # broken evaluator LIKELY
    "b7895e80", "58565672", "88fe4b2d",
    # live-web drift
    "a96b564e", "3e3fc409", "da46d875", "7ff48d5b", "deec51c9", "236833a3",
    "67890eb6", "d1acdb87",
    # infeasible task, agent said DONE
    "480bcfea", "2e6f678f", "a1af9f1c",
    # gate false positives
    "12086550", "2888b4e6",
}

meta = json.loads((OSW / "evaluation_examples/test_nogdrive.json").read_text())

# task_id -> domain for the universe
universe = {tid: dom for dom, tids in meta.items() for tid in tids}

# Pointer sonnet-4-6 step counts
steps = defaultdict(int)
with open(ROOT / "trials/trace-forensics/data/pointer_png_index.csv") as f:
    for r in csv.DictReader(f):
        if r["model"] == "sonnet-4-6":
            steps[r["task_id"]] = max(steps[r["task_id"]], int(r["step"]))

# Exclusion set from the classification CSV (anything not a genuine agent fail)
excluded = set()
with open(ROOT / "trials/trace-forensics/data/both_fail_classification.csv") as f:
    for r in csv.DictReader(f):
        if r["category"] != "AGENT_FAIL":
            excluded.add(r["task_id"])

# Infeasible-type tasks (evaluator func contains "infeasible")
infeasible = set()
for cfg in (OSW / "evaluation_examples/examples").glob("*/*.json"):
    try:
        d = json.loads(cfg.read_text())
    except Exception:
        continue
    fn = d.get("evaluator", {}).get("func")
    fns = fn if isinstance(fn, list) else [fn]
    if "infeasible" in fns:
        infeasible.add(cfg.stem)

def is_excluded(tid: str) -> bool:
    return (
        tid in excluded
        or tid in infeasible
        or any(tid.startswith(p) for p in FLAGGED_PREFIXES)
    )

# Candidate pools per stratum
pools = {"medium": defaultdict(list), "long": defaultdict(list)}
for tid, dom in universe.items():
    n = steps.get(tid, 0)
    if n == 0 or is_excluded(tid):
        continue
    if 9 <= n <= 15:
        pools["medium"][dom].append(tid)
    elif n >= 16:
        pools["long"][dom].append(tid)

for stratum in pools:
    for dom in pools[stratum]:
        # longer tasks first, then lexical for determinism
        pools[stratum][dom].sort(key=lambda t: (-steps[t], t))

def pick(pool: dict, n: int) -> list:
    """Round-robin across sorted domains until n tasks are picked."""
    out = []
    doms = sorted(pool)
    idx = {d: 0 for d in doms}
    while len(out) < n:
        progressed = False
        for d in doms:
            if len(out) >= n:
                break
            if idx[d] < len(pool[d]):
                out.append((d, pool[d][idx[d]]))
                idx[d] += 1
                progressed = True
        if not progressed:
            raise SystemExit(f"pool exhausted at {len(out)}/{n}")
    return out

picked = {"medium": pick(pools["medium"], 20), "long": pick(pools["long"], 20)}

OUT.mkdir(parents=True, exist_ok=True)

pack = defaultdict(list)
rows = []
for stratum in ("medium", "long"):
    for dom, tid in picked[stratum]:
        pack[dom].append(tid)
        rows.append({"domain": dom, "task_id": tid,
                     "pointer_steps": steps[tid], "stratum": stratum})

(OUT / "pack.json").write_text(json.dumps({d: pack[d] for d in sorted(pack)}, indent=2))

with open(OUT / "pack_detail.csv", "w", newline="") as f:
    w = csv.DictWriter(f, fieldnames=["domain", "task_id", "pointer_steps", "stratum"])
    w.writeheader()
    w.writerows(rows)

dom_counts = defaultdict(lambda: [0, 0])
for r in rows:
    dom_counts[r["domain"]][0 if r["stratum"] == "medium" else 1] += 1

md = ["# osworld-40-pack — shared 40-task OSWorld pack",
      "",
      "Built by trials/focuschain-ab/scripts/build_pack.py (deterministic; rerun to regenerate).",
      "",
      "- Universe: `evaluation_examples/test_nogdrive.json` (361 tasks)",
      "- Step basis: Pointer published sonnet-4-6 run, max PNG step per task",
      "  (`trials/trace-forensics/data/pointer_png_index.csv`)",
      "- Strata: **medium** = 9-15 Pointer steps (20 tasks), **long** = >=16 (20 tasks)",
      "- Excluded: all non-AGENT_FAIL tasks from",
      "  `trials/trace-forensics/data/both_fail_classification.csv` (broken evaluators,",
      "  live-web drift, infeasible-missed, gate false positives per",
      "  `findings/benchmark-integrity.md`) plus all 29 infeasible-type tasks",
      "  (evaluator func contains `infeasible`).",
      "- Within-stratum selection: round-robin over alphabetical domains; within a",
      "  domain, longest Pointer trajectory first.",
      "",
      "## Files",
      "- `pack.json` — `{domain: [task_ids]}`, drop-in for `--test_all_meta_path`",
      "- `pack_detail.csv` — domain, task_id, pointer_steps, stratum",
      "",
      "## Domain × stratum",
      "",
      "| domain | medium | long |",
      "|---|---:|---:|"]
for d in sorted(dom_counts):
    m, l = dom_counts[d]
    md.append(f"| {d} | {m} | {l} |")
md += ["", "## Tasks", "", "| domain | task_id | pointer_steps | stratum |", "|---|---|---:|---|"]
for r in rows:
    md.append(f"| {r['domain']} | {r['task_id']} | {r['pointer_steps']} | {r['stratum']} |")
(OUT / "pack.md").write_text("\n".join(md) + "\n")

print(f"medium pool: {sum(len(v) for v in pools['medium'].values())} candidates; "
      f"long pool: {sum(len(v) for v in pools['long'].values())}")
print(f"wrote {len(rows)} tasks to {OUT}")
for d in sorted(dom_counts):
    print(f"  {d}: medium={dom_counts[d][0]} long={dom_counts[d][1]}")
