#!/usr/bin/env python3
"""
Rebuild 002_final.srt (Myanmar) sau khi re-transcribe chunk 022.

Đọc toàn bộ chunk transcripts từ 002/transcripts/, merge + offset + clamp,
ghi lại 002_final.srt, 002_srt_map.json, 002_myanmar_lines.txt.

Đồng thời:
  - Ghi 002_rebuild_meta.json: {new_index, start, end, myanmar, chunk, old_index|null}
    → xác định segment nào cần dịch mới (chunk 022, old_index=null).
  - Ghi 002_new_to_translate.txt: danh sách các câu Myanmar mới cần dịch.

READ-ONLY đối với các file dịch (_vi.txt, batches) — chỉ GHI các file Myanmar + meta.

Usage:
    python3 rebuild_002.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"
TRANS_DIR = os.path.join(PROJ, PREFIX, "transcripts")
OUT_DIR = os.path.join(PROJ, PREFIX)
CHUNK_DURATION = 120.0
OVERLAP = 0.3
EFFECTIVE = CHUNK_DURATION - OVERLAP  # 119.7
GAP_MS = 30

TS_RE = re.compile(r'^(\d{1,2}):(\d{2}):(\d{2}),(\d{3})\s*-->\s*(\d{1,2}):(\d{2}):(\d{2}),(\d{3})$')


def to_ms(h, m, s, ms):
    return h * 3600000 + m * 60000 + s * 1000 + ms


def fmt_ms(ms):
    ms = int(round(ms))
    h = ms // 3600000; ms %= 3600000
    m = ms // 60000; ms %= 60000
    s = ms // 1000; ms %= 1000
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def parse_chunk_srt(path):
    """Trả về list [(start_ms, end_ms, text)] trong hệ giờ local của chunk."""
    segs = []
    for blk in open(path, encoding="utf-8").read().split("\n\n"):
        lines = blk.split("\n")
        if len(lines) < 3:
            continue
        m = TS_RE.match(lines[1].strip())
        if not m:
            continue
        g = [int(x) for x in m.groups()]
        s = to_ms(g[0], g[1], g[2], g[3])
        e = to_ms(g[4], g[5], g[6], g[7])
        text = "\n".join(lines[2:]).strip()
        if text:
            segs.append((s, e, text))
    return segs


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

    chunk_files = sorted(glob.glob(os.path.join(TRANS_DIR, f"{PREFIX}_clean_chunk_*_gemini.srt")))
    print(f"Found {len(chunk_files)} chunk SRTs")

    # ── Merge với offset ──
    all_segs = []  # [start_ms, end_ms, text, chunk_idx]
    for cf in chunk_files:
        m = re.search(r'chunk_(\d+)', cf)
        idx = int(m.group(1))
        offset_ms = int(round(idx * EFFECTIVE * 1000))
        for (s, e, text) in parse_chunk_srt(cf):
            all_segs.append([s + offset_ms, e + offset_ms, text, idx])

    all_segs.sort(key=lambda x: x[0])
    print(f"Total merged segments (raw): {len(all_segs)}")

    # ── Clamp overlap (end <= next start - 30ms, min 500ms) ──
    for i in range(len(all_segs)):
        if i < len(all_segs) - 1 and all_segs[i][1] > all_segs[i + 1][0] - GAP_MS:
            all_segs[i][1] = all_segs[i + 1][0] - GAP_MS
        if all_segs[i][1] <= all_segs[i][0]:
            all_segs[i][1] = all_segs[i][0] + 500

    n = len(all_segs)
    print(f"Total merged segments (final): {n}")

    # ── Phân đoạn cũ (collapsed chunk 022 = 1 segment tại index 612) ──
    # Chunk < 22 → giữ nguyên index; chunk 22 → NEW (cần dịch); chunk > 22 → shift +33
    new_records = []   # {start, end, myanmar}
    meta = []          # {new_index(0-based), chunk, old_index|null}
    for i, (s, e, text, cidx) in enumerate(all_segs):
        if cidx < 22:
            old_idx = i
        elif cidx == 22:
            old_idx = None
        else:
            # chunks > 22: old index = new index - 33 (vì 1 segment cũ → 34 segment mới)
            old_idx = i - 33
        new_records.append({"start": fmt_ms(s), "end": fmt_ms(e), "myanmar": text})
        meta.append({"new_index": i, "chunk": cidx, "old_index": old_idx})

    # ── Kiểm tra ổn định ──
    new_count = sum(1 for m in meta if m["old_index"] is None)
    print(f"Segments cần dịch mới (chunk 022): {new_count}")
    assert new_count == 34, f"kỳ vọng 34 segment mới, nhận {new_count}"

    if args.dry_run:
        print("[dry-run] không ghi file.")
        return

    # ── Ghi Myanmar files ──
    ts = time.strftime("%Y%m%d_%H%M%S")
    def bak(p):
        b = f"{p}.bak.{ts}"
        shutil.copy2(p, b)
        return b

    srt_path = os.path.join(OUT_DIR, f"{PREFIX}_final.srt")
    if os.path.exists(srt_path):
        bak(srt_path)
    with open(srt_path, "w", encoding="utf-8") as f:
        for i, r in enumerate(new_records, 1):
            f.write(f"{i}\n{r['start']} --> {r['end']}\n{r['myanmar']}\n\n")

    mp_path = os.path.join(OUT_DIR, f"{PREFIX}_srt_map.json")
    if os.path.exists(mp_path):
        bak(mp_path)
    records_1based = [{"index": i + 1, "start": r["start"], "end": r["end"],
                       "myanmar": r["myanmar"]} for i, r in enumerate(new_records)]
    with open(mp_path, "w", encoding="utf-8") as f:
        json.dump(records_1based, f, ensure_ascii=False, indent=2)

    ml_path = os.path.join(OUT_DIR, f"{PREFIX}_myanmar_lines.txt")
    if os.path.exists(ml_path):
        bak(ml_path)
    with open(ml_path, "w", encoding="utf-8") as f:
        for i, r in enumerate(new_records):
            f.write(f"[{i}] {r['myanmar']}\n")

    # ── Meta + danh sách cần dịch ──
    meta_path = os.path.join(OUT_DIR, f"{PREFIX}_rebuild_meta.json")
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

    new_txt_path = os.path.join(OUT_DIR, f"{PREFIX}_new_to_translate.txt")
    with open(new_txt_path, "w", encoding="utf-8") as f:
        for m in meta:
            if m["old_index"] is None:
                r = new_records[m["new_index"]]
                f.write(f"[{m['new_index']}] {r['myanmar']}\n")

    print(f"✅ 002_final.srt      ({n} cues)")
    print(f"✅ 002_srt_map.json")
    print(f"✅ 002_myanmar_lines.txt")
    print(f"✅ 002_rebuild_meta.json")
    print(f"✅ 002_new_to_translate.txt  ({new_count} câu cần dịch)")


if __name__ == "__main__":
    main()