#!/usr/bin/env python3
"""
Assemble 002_final_vi.srt + bilingual + batches sau khi re-transcribe chunk 022.

Dữ liệu vào:
  - 002_srt_map.json        (808 records Myanmar mới)
  - 002_rebuild_meta.json   (new_index → old_index mapping)
  - batches/002_batch_*_vi.txt  (bản dịch cũ, index 0-based cũ)
  - 002_new_vi.txt          (34 câu dịch mới cho chunk 022)

Output:
  - 002_final_vi.srt
  - 002_bilingual.txt
  - batches/002_batch_*_my.txt + _vi.txt (regenerate, 27 batches)

Usage:
    python3 assemble_002_vi.py [--dry-run]
"""
import argparse
import glob
import json
import os
import re
import shutil
import time

PROJ = "/home/tuan-nguyen/.openclaw/workspace/015-phu_de_video"
PREFIX = "002"
OUT = os.path.join(PROJ, PREFIX)
BATCH_SIZE = 30


def load_old_vi():
    vi = {}
    for path in sorted(glob.glob(os.path.join(OUT, "batches", f"{PREFIX}_batch_*_vi.txt"))):
        for line in open(path, encoding="utf-8"):
            line = line.strip()
            m = re.match(r'\[(\d+)\]\s*(.*)', line)
            if m:
                vi[int(m.group(1))] = m.group(2).strip()
    return vi


def load_new_vi():
    vi = {}
    path = os.path.join(OUT, f"{PREFIX}_new_vi.txt")
    for line in open(path, encoding="utf-8"):
        line = line.strip()
        m = re.match(r'\[(\d+)\]\s*(.*)', line)
        if m:
            vi[int(m.group(1))] = m.group(2).strip()
    return vi


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    records = json.load(open(os.path.join(OUT, f"{PREFIX}_srt_map.json"), encoding="utf-8"))
    meta = json.load(open(os.path.join(OUT, f"{PREFIX}_rebuild_meta.json"), encoding="utf-8"))
    assert len(records) == len(meta), f"srt_map {len(records)} != meta {len(meta)}"

    old_vi = load_old_vi()
    new_vi = load_new_vi()
    print(f"old vi loaded: {len(old_vi)} | new vi loaded: {len(new_vi)}")

    full_vi = []  # text theo new index 0-based
    missing = []
    for m in meta:
        new_idx = m["new_index"]
        old_idx = m["old_index"]
        if old_idx is None:
            t = new_vi.get(new_idx, "")
            if not t:
                missing.append(("new", new_idx))
        else:
            t = old_vi.get(old_idx, "")
            if not t:
                missing.append(("old", old_idx))
        full_vi.append(t)

    if missing:
        print(f"⚠️  THIẾU {len(missing)} bản dịch: {missing[:20]}")
        if not args.dry_run:
            print("→ ABORT: chưa ghi file vì còn thiếu bản dịch.")
            return 1

    if args.dry_run:
        print(f"[dry-run] {len(records)} cues, {len(full_vi)} vi, missing={len(missing)}")
        return 0

    ts = time.strftime("%Y%m%d_%H%M%S")

    # ── 002_final_vi.srt ──
    srt_path = os.path.join(OUT, f"{PREFIX}_final_vi.srt")
    if os.path.exists(srt_path):
        shutil.copy2(srt_path, f"{srt_path}.bak.{ts}")
    with open(srt_path, "w", encoding="utf-8") as f:
        for i, r in enumerate(records):
            f.write(f"{i+1}\n{r['start']} --> {r['end']}\n{full_vi[i]}\n\n")

    # ── 002_bilingual.txt ──
    bil_path = os.path.join(OUT, f"{PREFIX}_bilingual.txt")
    if os.path.exists(bil_path):
        shutil.copy2(bil_path, f"{bil_path}.bak.{ts}")
    with open(bil_path, "w", encoding="utf-8") as f:
        for i, r in enumerate(records):
            f.write(f"[{i}] {r['myanmar']}\n")
            f.write(f"     {full_vi[i]}\n")
            f.write("\n")

    # ── Regenerate batches ──
    n = (len(records) + BATCH_SIZE - 1) // BATCH_SIZE
    bdir = os.path.join(OUT, "batches")
    for b in range(n):
        s = b * BATCH_SIZE
        e = min(s + BATCH_SIZE, len(records))
        name = f"{PREFIX}_batch_{b+1:03d}"
        with open(os.path.join(bdir, f"{name}_my.txt"), "w", encoding="utf-8") as f:
            for i in range(s, e):
                f.write(f"[{i}] {records[i]['myanmar']}\n")
        with open(os.path.join(bdir, f"{name}_vi.txt"), "w", encoding="utf-8") as f:
            for i in range(s, e):
                f.write(f"[{i}] {full_vi[i]}\n")

    print(f"✅ 002_final_vi.srt  ({len(records)} cues)")
    print(f"✅ 002_bilingual.txt")
    print(f"✅ {n} batches regenerated (my + vi)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())