#!/usr/bin/env python3
"""Regenerate ``app/data/problem_ids.json`` — the judge endpoint's id whitelist.

    python scripts/gen_problem_ids.py \
        --judge-repo ../equational-theories-lean-stage2 [--out app/data/problem_ids.json]

``POST /v1/judge`` must reject an unknown ``problem_id`` before a job is ever
queued, and a regex cannot do it: ``^(normal|hard[123])_\\d{4}$`` happily matches
``hard1_0070``, which does not exist (hard1 stops at 0069). So the real id set
ships with the backend, generated here from the four canonical JSONLs in the
judge repo (normal 1000 + hard1 69 + hard2 200 + hard3 400 = 1669).

**IDS ONLY.** ERA2_DESIGN.md removes "the oracle surface on the service path:
the judge endpoint accepts problem_id only — never a statement, never an
answer". The canonical rows carry both ``equation1``/``equation2`` and the
``answer`` bool, and this backend is the PUBLIC face of the collab, so neither
may land in its image: this writes the ids and provenance and nothing else
(~20 KB). The statements stay where the judge already needs them — inside the
private eval Space's own checkout of the judge repo.

Provenance (repo path, per-file sha256, count, generated_at) is recorded so a
later reader can prove which checkout the ids came from; ``--check`` verifies
the committed file still matches the repo without rewriting it.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import sys
from datetime import datetime, timezone
from pathlib import Path


SETS = {"normal": 1000, "hard1": 69, "hard2": 200, "hard3": 400}
EXPECTED_TOTAL = 1669
DEFAULT_OUT = Path(__file__).resolve().parent.parent / "app" / "data" / "problem_ids.json"


def build(judge_repo: Path) -> dict:
    problems_dir = judge_repo / "examples" / "problems"
    ids: list[str] = []
    sources: dict[str, dict] = {}
    for name, expected in SETS.items():
        f = problems_dir / f"{name}.jsonl"
        raw = f.read_bytes()
        set_ids = [
            json.loads(line)["id"]
            for line in raw.decode("utf-8").splitlines()
            if line.strip()
        ]
        if len(set_ids) != expected:
            raise SystemExit(f"{f}: expected {expected} rows, found {len(set_ids)}")
        sources[f"{name}.jsonl"] = {
            "count": len(set_ids),
            "sha256": hashlib.sha256(raw).hexdigest(),
        }
        ids.extend(set_ids)

    if len(set(ids)) != len(ids):
        raise SystemExit("duplicate problem ids across the canonical sets")
    if len(ids) != EXPECTED_TOTAL:
        raise SystemExit(f"expected {EXPECTED_TOTAL} canonical ids, built {len(ids)}")

    return {
        "count": len(ids),
        "provenance": {
            "judge_repo": str(judge_repo),
            "files": sources,
            "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
            "note": "ids only — no statements, no answers (ERA2_DESIGN.md)",
        },
        "ids": sorted(ids),
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--judge-repo", required=True, type=Path,
                    help="checkout of equational-theories-lean-stage2")
    ap.add_argument("--out", type=Path, default=DEFAULT_OUT)
    ap.add_argument("--check", action="store_true",
                    help="compare the existing file's ids with the repo; write nothing")
    args = ap.parse_args()

    doc = build(args.judge_repo.resolve())
    if args.check:
        current = json.loads(args.out.read_text())
        if current.get("ids") != doc["ids"]:
            print(f"STALE: {args.out} ids differ from {args.judge_repo}")
            return 1
        print(f"ok: {args.out} matches {args.judge_repo} ({doc['count']} ids)")
        return 0

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(doc, indent=1, sort_keys=False) + "\n")
    print(f"wrote {args.out} ({doc['count']} ids, {args.out.stat().st_size} bytes)")
    return 0


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