"""A fixed test subset for the LLM comparisons: N rooms of test.jsonl, stratified by source in proportion to the test
set (largest remainder; a source too small for one room gets none), each source's rooms drawn with one seeded RNG.
Rooms over the model capacity (more than --max-objects objects) are left out before drawing, so every method can
answer every room. Writes rows.jsonl (test order) and MANIFEST.json (seed, counts, sha256, scene ids).

    python fastfill/v2/ops/make_test300.py --test data/v4fixed/test.jsonl --output data/test300 [--rows 300 --seed 0]
"""
from __future__ import annotations

import argparse
from collections import defaultdict
import hashlib
import json
from pathlib import Path
import random


def allocate(sizes: dict, total: int) -> dict:
    """Largest-remainder apportionment of ``total`` over ``sizes`` (ties broken by name)."""
    whole = sum(sizes.values())
    quota = {k: total * v / whole for k, v in sizes.items()}
    share = {k: int(q) for k, q in quota.items()}
    for k in sorted(quota, key=lambda k: (-(quota[k] - share[k]), k))[:total - sum(share.values())]:
        share[k] += 1
    return share


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--test", type=Path, required=True)
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--rows", type=int, default=300)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--max-objects", type=int, default=256)
    a = p.parse_args(argv)
    lines = [l for l in a.test.read_text().splitlines() if l.strip()]
    by_source, over = defaultdict(list), 0
    for index, line in enumerate(lines):
        row = json.loads(line)
        if len(row["condition"]["objects"]) > a.max_objects:
            over += 1
            continue
        by_source[row["provenance"]["source"]].append(index)
    share, rng, chosen = allocate({k: len(v) for k, v in by_source.items()}, a.rows), random.Random(a.seed), []
    for source in sorted(by_source):
        chosen += rng.sample(by_source[source], share[source])
    chosen.sort()
    a.output.mkdir(parents=True, exist_ok=True)
    rows = "".join(lines[i] + "\n" for i in chosen)
    (a.output / "rows.jsonl").write_text(rows)
    manifest = {"rows": len(chosen), "seed": a.seed, "max_objects": a.max_objects, "over_capacity_excluded": over,
                "test_file": a.test.name, "test_rows": len(lines), "test_sha256": hashlib.sha256(a.test.read_bytes()).hexdigest(),
                "rows_sha256": hashlib.sha256(rows.encode()).hexdigest(),
                "per_source": {k: {"test": len(by_source[k]), "chosen": share[k]} for k in sorted(by_source)},
                "test_line_indices": chosen,
                "scene_ids": [json.loads(lines[i])["provenance"].get("scene_id") for i in chosen]}
    (a.output / "MANIFEST.json").write_text(json.dumps(manifest, indent=1) + "\n")
    print(json.dumps({k: manifest[k] for k in ("rows", "over_capacity_excluded", "rows_sha256")}),
          {k: v["chosen"] for k, v in manifest["per_source"].items()})


if __name__ == "__main__":
    main()