#!/usr/bin/env python3
"""Build the BrowseComp-Plus structures bundle from an AML response JSONL.

Reads the per-question structure-generation responses (from the
``information-scaffolds`` derive-structure pipeline) and writes one
combined JSONL into ``backend/data/structures.jsonl``. The Flask backend
loads this file at startup and serves it via /api/browsecomp/structures/.

Single-file output (not per-qid shards) keeps the BrowseComp Space:
  * inside the LFS rule already in .gitattributes (``backend/data/*.jsonl``),
  * cheap to ship inside the Docker image (one COPY, one file open).

Usage:
    python scripts/build_structures.py /path/to/response.jsonl
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys

DATASET = "browsecomp_plus"
DEFAULT_OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "backend",
    "data",
    "structures.jsonl",
)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("response", help="Path to the AML response JSONL file")
    ap.add_argument(
        "--out",
        default=DEFAULT_OUT,
        help="Output JSONL path (default: backend/data/structures.jsonl)",
    )
    ap.add_argument(
        "--include-empty",
        action="store_true",
        help="Keep rows whose answer is empty/whitespace (default: skip — these are "
             "model timeouts or content-filter rejections with no useful structure).",
    )
    args = ap.parse_args()

    records = []
    skipped_empty = 0
    with open(args.response) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            if obj.get("dataset") != DATASET:
                continue
            answer = obj.get("answer", "") or ""
            if not args.include_empty and not answer.strip():
                skipped_empty += 1
                continue
            prompt_doc_ids = re.findall(r"\[Doc \d+\] id=(\S+)", obj.get("user_prompt", ""))
            records.append({
                "qid": str(obj["qid"]),
                "question": obj["question"],
                "model": obj.get("model"),
                "reasoning_effort": obj.get("reasoning_effort"),
                "max_completion_tokens": obj.get("max_completion_tokens"),
                "available_num_docs": obj.get("available_num_docs"),
                "prompt_num_docs": obj.get("prompt_num_docs"),
                "system_prompt_file": obj.get("system_prompt_file"),
                "prompt_doc_ids": prompt_doc_ids,
                "usage": obj.get("usage"),
                "latency_ms": obj.get("latency_ms"),
                "finish_reason": obj.get("finish_reason"),
                "answer": answer,
            })

    records.sort(key=lambda r: r["qid"])
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    with open(args.out, "w") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print(
        f"wrote {len(records)} {DATASET} structure record(s) to {args.out}"
        + (f" (skipped {skipped_empty} empty-answer row(s))" if skipped_empty else ""),
        file=sys.stderr,
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())