"""Write the public dataset files from data/final into release/dataset/data with a stable schema.

  .venv/bin/python -I scripts/export_release_data.py [--run-a data/final_random/sft_verified.jsonl]

  sft.parquet         one verified-correct Beam trace per prompt (shortest correct), messages form
  traces_all.parquet  every generation, correct or not (held-out eval prompts are never included)

Everything in the dataset repo is Parquet: the Hub picks one file reader per dataset, so formats can't be mixed.
  run_a_ids.txt     prompt ids in the 52,660-example snapshot Run A was trained on
"""

import argparse
import json
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

ROOT = Path(__file__).resolve().parent.parent
TEACHER = {"beam-api": "api", "beam-nvfp4-selfhost-sglang": "selfhost-nvfp4", "beam-fp8-selfhost-sglang": "selfhost-fp8"}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--final", default=str(ROOT / "data" / "final"))
    ap.add_argument("--run-a", default=str(ROOT / "data" / "final_random" / "sft_verified.jsonl"),
                    help="any SFT build from the Run A snapshot (same prompt ids)")
    ap.add_argument("--out", default=str(ROOT / "release" / "dataset" / "data"))
    args = ap.parse_args()
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    eval_ids = {l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()}

    rows = []
    for line in open(Path(args.final) / "sft_verified.jsonl"):
        r = json.loads(line)
        rows.append({"id": r["id"], "domain": r["domain"], "source": r["source"], "difficulty": r["difficulty"],
                     "teacher": TEACHER[r["teacher"]], "reasoning_effort": r["effort"],
                     "messages": r["messages"], "gold": r["gold"],
                     "n_samples": r["n_samples"], "n_correct": r["n_correct"]})
    pq.write_table(pa.Table.from_pylist(rows), out / "sft.parquet", compression="zstd")
    print(f"sft.parquet: {len(rows)}")

    rows = []
    for line in open(Path(args.final) / "traces_all.jsonl"):
        r = json.loads(line)
        assert r["id"] not in eval_ids, f"held-out eval prompt {r['id']} in traces"
        rows.append({
            "id": r["id"], "domain": r["domain"], "source": r["source"], "topic": r.get("topic"),
            "difficulty": r["difficulty"],
            "teacher": TEACHER[r.get("teacher", "beam-api")], "reasoning_effort": r["effort"],
            "prompt": r["prompt"], "reasoning": r["reasoning"], "content": r["content"],
            "gold": r["gold"], "answer_type": r["answer_type"], "pred": r["pred"], "correct": r["correct"],
            "finish_reason": r["finish_reason"],
            # SGLang n>1 responses report only total usage, so self-hosted counts are per-request estimates
            "completion_tokens": r["usage"].get("completion_tokens"),
        })
    pq.write_table(pa.Table.from_pylist(rows), out / "traces_all.parquet", compression="zstd")
    print(f"traces_all.parquet: {len(rows)}")

    ids = sorted({json.loads(l)["id"] for l in open(args.run_a)})
    (out.parent / "run_a_ids.txt").write_text("\n".join(ids) + "\n")
    print(f"run_a_ids.txt: {len(ids)}")


if __name__ == "__main__":
    main()