import json, csv, argparse
from pathlib import Path
from collections import defaultdict

def load_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                yield json.loads(line)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--raw_dir", required=True)  # eval/runs/<EVAL_ID>/raw
    ap.add_argument("--out_dir")                 # default: ../summary
    args = ap.parse_args()

    raw_dir = Path(args.raw_dir)
    out_dir = Path(args.out_dir) if args.out_dir else raw_dir.parent / "summary"
    out_dir.mkdir(parents=True, exist_ok=True)

    runs = defaultdict(lambda: {
        "eval_id": None,
        "git_sha": None,
        "model_ids": set(),
        "scenarios": set(),
        "files": 0,
        "turns": 0,
        "errors": 0,
        "start_ts": None,
        "end_ts": None,
    })

    for jf in raw_dir.glob("*.jsonl"):
        # infer run_id from filename suffix _xxxxxx.jsonl
        name = jf.stem
        run_id = name.rsplit("_", 1)[-1]
        runs[run_id]["files"] += 1

        for row in load_jsonl(jf):
            rid = row.get("run_id", run_id)
            if rid != run_id:
                run_id = rid  # prefer payload
            r = runs[run_id]

            r["turns"] += 1
            if row.get("status") != "ok":
                r["errors"] += 1

            r["eval_id"] = r["eval_id"] or row.get("eval_id")
            r["git_sha"] = r["git_sha"] or row.get("git_sha")
            if row.get("model_id"):
                r["model_ids"].add(row["model_id"])

            scn = row.get("scenario_id")
            if scn:
                r["scenarios"].add(scn)

            st = row.get("start_ts")
            et = row.get("end_ts")
            if st and (r["start_ts"] is None or st < r["start_ts"]):
                r["start_ts"] = st
            if et and (r["end_ts"] is None or et > r["end_ts"]):
                r["end_ts"] = et

    # write JSON
    json_out = {}
    for rid, d in runs.items():
        json_out[rid] = {
            "eval_id": d["eval_id"],
            "git_sha": d["git_sha"],
            "model_ids": sorted(d["model_ids"]),
            "scenarios": sorted(d["scenarios"]),
            "files": d["files"],
            "turns": d["turns"],
            "errors": d["errors"],
            "start_ts": d["start_ts"],
            "end_ts": d["end_ts"],
        }
    (out_dir / "runs_index.json").write_text(json.dumps(json_out, indent=2), encoding="utf-8")

    # write CSV
    with open(out_dir / "runs_index.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["run_id","eval_id","git_sha","model_ids","files","scenarios","turns","errors","start_ts","end_ts"])
        for rid, d in sorted(runs.items()):
            w.writerow([
                rid,
                d["eval_id"],
                d["git_sha"],
                "|".join(sorted(d["model_ids"])) if d["model_ids"] else "",
                d["files"],
                len(d["scenarios"]),
                d["turns"],
                d["errors"],
                d["start_ts"],
                d["end_ts"],
            ])

    print(f"✔ wrote {out_dir/'runs_index.json'}")
    print(f"✔ wrote {out_dir/'runs_index.csv'}")

if __name__ == "__main__":
    main()