#!/usr/bin/env python3
"""把對白切成給標註 sub-agent 的批次檔。

用法：
  make_batches.py <dialogues.jsonl> <輸出目錄> --pilot N   # 試跑：分層抽樣 N 句
  make_batches.py <dialogues.jsonl> <輸出目錄> --all       # 全量：規則未判定或中信心的對白

每批一個 jsonl，一行一句，附前後各 2 句對白當輪流發言的線索。
批次檔不含規則判定的說話者，讓標註獨立，之後才能拿來和規則比對。
"""
import argparse
import json
import random
from collections import defaultdict
from pathlib import Path


def short(rec, n=60):
    t = rec["text"]
    return {"id": rec["id"], "text": t if len(t) <= n else t[:n] + "…"}


def with_neighbors(recs):
    by_chapter = defaultdict(list)
    for r in recs:
        by_chapter[(r["volume"], r["chapter"])].append(r)
    index = {}
    for rows in by_chapter.values():
        for i, r in enumerate(rows):
            index[r["id"]] = {"prev": [short(x) for x in rows[max(0, i - 2):i]], "next": [short(x) for x in rows[i + 1:i + 3]]}
    return index


def pilot_sample(recs, n, seed):
    rng = random.Random(seed)
    unlabeled = [r for r in recs if not r["speaker"]]
    medium = [r for r in recs if r["confidence"] == "medium"]
    high = [r for r in recs if r["confidence"] == "high"]
    take = [(unlabeled, n // 2), (medium, n // 4), (high, n - n // 2 - n // 4)]
    out = []
    for pool, k in take:
        out += rng.sample(pool, min(k, len(pool)))
    rng.shuffle(out)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dialogues")
    ap.add_argument("outdir")
    ap.add_argument("--pilot", type=int)
    ap.add_argument("--all", action="store_true")
    ap.add_argument("--batch-size", type=int, default=25)
    ap.add_argument("--seed", type=int, default=7)
    a = ap.parse_args()
    if bool(a.pilot) == a.all:
        ap.error("--pilot N 與 --all 擇一")
    recs = [json.loads(line) for line in Path(a.dialogues).read_text(encoding="utf-8").splitlines()]
    nb = with_neighbors(recs)
    chosen = pilot_sample(recs, a.pilot, a.seed) if a.pilot else [r for r in recs if r["confidence"] != "high"]
    out = Path(a.outdir)
    out.mkdir(parents=True, exist_ok=True)
    for i in range(0, len(chosen), a.batch_size):
        path = out / f"batch-{i // a.batch_size + 1:03d}.jsonl"
        with path.open("w", encoding="utf-8") as fh:
            for r in chosen[i:i + a.batch_size]:
                item = {k: r[k] for k in ("id", "text", "before", "after", "prev_par", "next_par", "candidates", "suspect_duplicate")}
                item.update(nb[r["id"]])
                fh.write(json.dumps(item, ensure_ascii=False) + "\n")
    print(f"{len(chosen)} 句 -> {(len(chosen) + a.batch_size - 1) // a.batch_size} 批 -> {out}")


if __name__ == "__main__":
    main()