"""Phase C — assign table dumps to extraction batches (balanced by dump size) and write the batch manifests.
Usage: python scripts/extraction/make_batches.py --n 8 [--only screened]  -> data/extracted/batches/batch_<i>.txt (list of record_ids)."""
import argparse, json
from pathlib import Path
import pandas as pd

ROOT = Path(__file__).resolve().parents[2]; EX = ROOT / "data" / "extracted"


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--n", type=int, default=8); ap.add_argument("--max_chars", type=int, default=60000); ap.add_argument("--start", type=int, default=1); a = ap.parse_args()
    idx = pd.read_csv(EX / "pdf_index.csv"); idx = idx[(~idx.scanned) & (idx.n_candidate_pages > 0)]
    done = {p.stem for p in (EX / "json").glob("*.json")} if (EX / "json").exists() else set()
    assigned = set()
    for b in (EX / "batches").glob("batch_*.txt"):
        assigned |= set(b.read_text().split())
    idx = idx[~idx.record_id.astype(str).isin(done | assigned)]
    # screening decisions (if available): keep includes/unclear
    dec_files = sorted((ROOT / "data" / "metadata" / "screening_batches").glob("decisions_*.csv"))
    if dec_files:
        dec = pd.concat(pd.read_csv(f) for f in dec_files)
        keep = set(dec[dec.reason_code.eq("ELIGIBLE_HEDONIC") | dec.decision.eq("unclear")].record_id.astype(str))
        idx = idx[idx.record_id.astype(str).isin(keep) | ~idx.record_id.astype(str).isin(set(dec.record_id.astype(str)))]
    idx = idx.sort_values("dump_chars", ascending=False)
    batches = [[] for _ in range(a.n)]; sizes = [0] * a.n
    for r in idx.itertuples():
        j = sizes.index(min(sizes)); batches[j].append(r.record_id); sizes[j] += min(r.dump_chars, a.max_chars)
    (EX / "batches").mkdir(exist_ok=True)
    for i, b in enumerate(batches, a.start):
        (EX / "batches" / f"batch_{i}.txt").write_text("\n".join(map(str, b)))
    print(len(idx), "dumps in", a.n, "batches; chars per batch:", sizes)


if __name__ == "__main__":
    main()