#!/usr/bin/env python3
"""挑出需要用新版指示重標的 none 對白。

用法：make_relabel.py <labeling/full 目錄> <輸出目錄> --max-batch N

早期批次用的指示把「嵌在敘述中的片段」一律標 none，會漏掉確實是某人說出口的話。
這裡從批次 1..N 的標註結果挑出 speaker=none 且看起來像完整句子的對白（長度 ≥ 6 字，
或含標點），重新切成批次檔，格式與 make_batches.py 相同。
"""
import argparse
import json
import re
from pathlib import Path

SENTENCE_HINT = re.compile(r"[，。！？…、]")


def looks_spoken(text):
    return len(text) >= 6 or bool(SENTENCE_HINT.search(text))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("full_dir")
    ap.add_argument("outdir")
    ap.add_argument("--max-batch", type=int, required=True)
    ap.add_argument("--batch-size", type=int, default=50)
    a = ap.parse_args()
    full = Path(a.full_dir)
    picked = []
    for n in range(1, a.max_batch + 1):
        name = f"batch-{n:03d}.jsonl"
        src = {json.loads(l)["id"]: json.loads(l) for l in (full / "in" / name).read_text(encoding="utf-8").splitlines()}
        for line in (full / "out" / name).read_text(encoding="utf-8").splitlines():
            lb = json.loads(line)
            if lb["speaker"] == "none" and looks_spoken(src[lb["id"]]["text"]):
                picked.append(src[lb["id"]])
    out = Path(a.outdir)
    out.mkdir(parents=True, exist_ok=True)
    for i in range(0, len(picked), a.batch_size):
        path = out / f"batch-{i // a.batch_size + 1:03d}.jsonl"
        path.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in picked[i:i + a.batch_size]), encoding="utf-8")
    print(f"{len(picked)} 句 -> {(len(picked) + a.batch_size - 1) // a.batch_size} 批 -> {out}")


if __name__ == "__main__":
    main()