#!/usr/bin/env python3
"""Build open-wikitable `eval_structures/` shards from the unified eval bundle.

Mirrors monaco-benchmark-viewer/scripts/build_eval_structures.py (same I/O
shape: per-qid JSON shard under ``eval_structures/records/<qid>.json`` plus a
single ``index.json`` summary).

Reads ``test_with_structures.unified.jsonl`` (which is the standard unified
question schema with extra ``docs[].id == "structures/<subtype>/<title>..."``
synthetic typed-structure docs emitted by data_creation/). Keeps only those
structures docs and emits one shard per qid; the viewer JS derives subtype/
title from each doc id and renders per subtype.

Usage:
    python scripts/build_eval_structures.py
    python scripts/build_eval_structures.py --src /path/to/test_with_structures.unified.jsonl
"""
from __future__ import annotations

import argparse
import json
import os
import sys

DATASET = "wiki_opentable"
HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_OUT = os.path.join(os.path.dirname(HERE), "eval_structures")
DEFAULT_DATA_ROOT = os.environ.get("DATA_ROOT", "/mnt/ramdisk/blobstore/timchen0618/data")
DEFAULT_SRC = os.path.join(DEFAULT_DATA_ROOT, "eval", DATASET, "unified", "test_with_structures.unified.jsonl")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=DEFAULT_SRC, help=f"Source unified jsonl (default: {DEFAULT_SRC})")
    ap.add_argument("--out", default=DEFAULT_OUT, help="Output eval_structures/ dir (default: ./eval_structures)")
    ap.add_argument(
        "--include-empty",
        action="store_true",
        help="Keep records with zero structure docs (default: skip them).",
    )
    args = ap.parse_args()

    if not os.path.exists(args.src):
        print(f"error: source file not found: {args.src}", file=sys.stderr)
        return 2

    records = []
    skipped_empty = 0
    with open(args.src) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            docs = obj.get("docs") or []
            structures = [
                {"id": d.get("id"), "contents": d.get("contents", "")}
                for d in docs
                if (d.get("id") or "").startswith("structures/")
            ]
            if not structures and not args.include_empty:
                skipped_empty += 1
                continue
            records.append({
                "qid": str(obj["qid"]),
                "question": obj.get("question", ""),
                "answers": obj.get("answers") or [],
                "dataset_origin": obj.get("dataset_origin"),
                "original_table_id": obj.get("original_table_id"),
                "structures": structures,
            })

    records.sort(key=lambda r: r["qid"])
    rec_dir = os.path.join(args.out, "records")
    os.makedirs(rec_dir, exist_ok=True)
    for stale in os.listdir(rec_dir):
        if stale.endswith(".json"):
            os.remove(os.path.join(rec_dir, stale))

    with open(os.path.join(args.out, "index.json"), "w") as f:
        json.dump(
            [
                {
                    "qid": r["qid"],
                    "question": r["question"],
                    "dataset_origin": r["dataset_origin"],
                    "n_structures": len(r["structures"]),
                }
                for r in records
            ],
            f,
            ensure_ascii=False,
        )
    for r in records:
        with open(os.path.join(rec_dir, f"{r['qid']}.json"), "w") as f:
            json.dump(r, f, ensure_ascii=False)

    msg = f"wrote {len(records)} {DATASET} eval-structure record(s) to {args.out}"
    if skipped_empty:
        msg += f" (skipped {skipped_empty} record(s) with no structure docs)"
    print(msg, file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())