#!/usr/bin/env python3
"""Build BrowseComp-Plus backend/data/test_with_structures.unified.jsonl
from the unified eval bundle.

Reads the unified test_with_structures jsonl, keeps only docs whose `id`
starts with `structures/` (= synthetic typed-structure docs emitted by
data_creation/), and writes a single jsonl with one record per qid (matching
the shape the unified loader uses). The Flask backend then serves it from
/api/browsecomp/eval_structures/.

Unlike the static viewers, BrowseComp does *not* shard per qid — it caches
the whole file in memory at module load time (same pattern as the existing
unified + structures endpoints), and the file is committed under
backend/data/*.jsonl (LFS-tracked via the existing rule in .gitattributes).

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 = "browsecomp_plus"
DEFAULT_OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "backend", "data", "test_with_structures.unified.jsonl",
)
DEFAULT_DATA_ROOT = os.environ.get("DATA_ROOT", "/tmp/blobfuse2/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=f"Output jsonl (default: {DEFAULT_OUT})")
    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

    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    n_kept = 0
    n_skipped = 0
    with open(args.src) as f_in, open(args.out, "w") as f_out:
        for line in f_in:
            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:
                n_skipped += 1
                continue
            out_rec = {
                "qid": str(obj["qid"]),
                "question": obj.get("question", ""),
                "structures": structures,
            }
            f_out.write(json.dumps(out_rec, ensure_ascii=False) + "\n")
            n_kept += 1
    msg = f"wrote {n_kept} {DATASET} eval-structure record(s) to {args.out}"
    if n_skipped:
        msg += f" (skipped {n_skipped} record(s) with no structure docs)"
    print(msg, file=sys.stderr)
    return 0


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