"""Write a JSON record of a train/eval split: ids and per-scenario counts per file, plus which scenarios are missing.

usage: python scripts/record_split.py out.json name1=file1.jsonl name2=file2.jsonl ...
"""
import json, sys
from collections import Counter

out, specs = sys.argv[1], sys.argv[2:]
res = {}
for s in specs:
    name, path = s.split("=", 1)
    with open(path, encoding="utf-8") as f:
        rows = [json.loads(l) for l in f if l.strip()]
    sc = Counter((r.get("meta") or {}).get("scenario") or "?" for r in rows)
    res[name] = {"path": path, "n": len(rows), "n_with_box_questions": sum(any(q.get("type") == "box" for q in r["questions"].values()) for r in rows),
                 "scenarios": dict(sorted(sc.items())), "ids": [r["id"] for r in rows]}
allsc = set().union(*(set(v["scenarios"]) for v in res.values()))
for v in res.values():
    v["missing_scenarios"] = sorted(allsc - set(v["scenarios"]))
with open(out, "w", encoding="utf-8", newline="\n") as f:
    json.dump(res, f, indent=1, ensure_ascii=False)
print(json.dumps({k: {"n": v["n"], "box": v["n_with_box_questions"], "n_scen": len(v["scenarios"]), "missing": v["missing_scenarios"]} for k, v in res.items()}))