#!/usr/bin/env python3
"""Build per-qid record shards for the MoNaCo dev viewer's Dataset tab.

Reads the raw monaco_dev split
    $DATA_ROOT/eval/monaco_dev/raw/test_with_chunks.jsonl
and emits per-qid record files under `<out-dir>/records/<qid>.json` plus an
`<out-dir>/index.json` sidebar index.

Record layout (matches viewer.js `renderDatasetRecord`):
{
  "dataset": "monaco_dev",
  "qid": "0",
  "ex_num": 0,
  "question": "...",
  "answers": ["a", "b", "c"],
  "docs": [
    {"id": "...", "title": "...", "section": "...", "sub_question_text": "...", "text": "..."},
    ...
  ]
}

Index row:
{"qid": "0", "question": "...", "n_answers": 3, "n_docs": 103}

Notes:
- monaco raw records only carry supporting docs (every `docs[i].is_supporting`
  is True). The full 6184-doc BM25-padded pool is in `raw/corpus.jsonl` and
  is not per-qid, so we surface just the supporting docs here.
- qid = str(ex_num) — mirrors the unified schema's qid convention.
"""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument(
        "--in-raw",
        default="/mnt/ramdisk/blobstore/timchen0618/data/eval/monaco_dev/raw/test_with_chunks.jsonl",
        help="path to monaco_dev raw test file",
    )
    ap.add_argument(
        "--out-dir",
        default="/home/azureuser/projects/data_visualizer/monaco-dev-viewer",
        help="viewer repo root — writes index.json and records/*.json inside it",
    )
    ap.add_argument("--pretty", action="store_true", help="pretty-print per-qid JSON")
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    records_dir = out_dir / "records"
    if records_dir.exists():
        shutil.rmtree(records_dir)
    records_dir.mkdir(parents=True)

    index: list[dict] = []
    n = 0
    with open(args.in_raw) as f:
        for line in f:
            r = json.loads(line)
            qid = str(r["ex_num"])
            docs = [
                {
                    "id": d.get("id"),
                    "title": d.get("title"),
                    "section": d.get("section"),
                    "sub_question_text": d.get("sub_question_text"),
                    "text": d.get("text") or d.get("contents") or "",
                }
                for d in (r.get("docs") or [])
            ]
            rec = {
                "dataset": "monaco_dev",
                "qid": qid,
                "ex_num": r["ex_num"],
                "question": r["question"],
                "answers": r.get("answers") or [],
                "docs": docs,
            }
            out_path = records_dir / f"{qid}.json"
            with open(out_path, "w") as g:
                if args.pretty:
                    json.dump(rec, g, indent=2, ensure_ascii=False)
                else:
                    json.dump(rec, g, ensure_ascii=False)
            index.append({
                "qid": qid,
                "question": r["question"],
                "n_answers": len(rec["answers"]),
                "n_docs": len(docs),
            })
            n += 1

    with open(out_dir / "index.json", "w") as g:
        json.dump(index, g, indent=2, ensure_ascii=False)

    print(f"wrote {n} records -> {records_dir}")
    print(f"wrote index      -> {out_dir/'index.json'}")


if __name__ == "__main__":
    main()