#!/usr/bin/env python3
"""
Fix 001: gộp cue 644 (fragment 170ms, đứt giữa từ 'နှုတ်အပြတ်') vào cue 645.

Nguyên nhân: từ 'နှုတ်အပြတ်' bị cắt đôi ở ranh giới chunk (00:59:50,970 ≈ 3591s).

Merge cue 644+645 (0-based 643+644) → 1 cue. Renumber mọi cue sau (shift -1).
Ghi lại: 001_final.srt, 001_final_vi.srt, 001_srt_map.json,
        001_myanmar_lines.txt, 001_bilingual.txt + regenerate batches.

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

PROJ = "/home/tuan-nguyen/.openclaw/workspace/015-phu_de_video"
PREFIX = "001"
OUT = os.path.join(PROJ, PREFIX)
MERGE_I0 = 643  # 0-based index của cue 644 (1-based)
BATCH_SIZE = 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(t):
    h, m, s, ms = [int(x) for x in re.split(r'[:,]', t)]
    return h * 3600000 + m * 60000 + s * 1000 + ms


def parse_srt(path):
    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(f"{g[0]}:{g[1]}:{g[2]}:{g[3]}")
        e = to_ms(f"{g[4]}:{g[5]}:{g[6]}:{g[7]}")
        text = "\n".join(lines[2:]).strip()
        segs.append({"start_ms": s, "end_ms": e, "text": text})
    return segs


def fmt_ms(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 main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    my = parse_srt(os.path.join(OUT, f"{PREFIX}_final.srt"))
    vi = parse_srt(os.path.join(OUT, f"{PREFIX}_final_vi.srt"))
    assert len(my) == len(vi), f"my={len(my)} vi={len(vi)}"

    a, b = my[MERGE_I0], my[MERGE_I0 + 1]
    print(f"Cue {MERGE_I0+1} (0-based {MERGE_I0}): dur={(a['end_ms']-a['start_ms'])}ms  "
          f"'{a['text'][-12:]}...'")
    print(f"Cue {MERGE_I0+2} (0-based {MERGE_I0+1}): starts '{b['text'][:12]}...'")
    assert a["end_ms"] - a["start_ms"] < 300, "cue 644 không còn là fragment ngắn — kiểm tra lại"

    # Merge Myanmar (nối trực tiếp: 'နှုတ်' + 'အပြတ်...' → 'နှုတ်အပြတ်...')
    my_merged_text = a["text"].rstrip() + b["text"].lstrip()
    # Merge Vietnamese (nối với dấu cách)
    vi_merged_text = (vi[MERGE_I0]["text"].rstrip() + " " + vi[MERGE_I0 + 1]["text"].lstrip())

    new_my = my[:MERGE_I0] + [{"start_ms": a["start_ms"], "end_ms": b["end_ms"],
                               "text": my_merged_text}] + my[MERGE_I0 + 2:]
    new_vi = vi[:MERGE_I0] + [{"start_ms": a["start_ms"], "end_ms": b["end_ms"],
                               "text": vi_merged_text}] + vi[MERGE_I0 + 2:]

    print(f"721 → {len(new_my)} cues")
    print(f"merged: {fmt_ms(a['start_ms'])} --> {fmt_ms(b['end_ms'])}")
    print(f"  MY: {my_merged_text[:60]}...")
    print(f"  VI: {vi_merged_text[:60]}...")

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

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

    # final.srt (Myanmar)
    p = os.path.join(OUT, f"{PREFIX}_final.srt"); bak(p)
    with open(p, "w", encoding="utf-8") as f:
        for i, s in enumerate(new_my, 1):
            f.write(f"{i}\n{fmt_ms(s['start_ms'])} --> {fmt_ms(s['end_ms'])}\n{s['text']}\n\n")

    # final_vi.srt
    p = os.path.join(OUT, f"{PREFIX}_final_vi.srt"); bak(p)
    with open(p, "w", encoding="utf-8") as f:
        for i, s in enumerate(new_vi, 1):
            f.write(f"{i}\n{fmt_ms(s['start_ms'])} --> {fmt_ms(s['end_ms'])}\n{s['text']}\n\n")

    # srt_map.json
    p = os.path.join(OUT, f"{PREFIX}_srt_map.json"); bak(p)
    recs = [{"index": i + 1, "start": fmt_ms(s["start_ms"]), "end": fmt_ms(s["end_ms"]),
             "myanmar": s["text"]} for i, s in enumerate(new_my)]
    with open(p, "w", encoding="utf-8") as f:
        json.dump(recs, f, ensure_ascii=False, indent=2)

    # myanmar_lines.txt
    p = os.path.join(OUT, f"{PREFIX}_myanmar_lines.txt"); bak(p)
    with open(p, "w", encoding="utf-8") as f:
        for i, s in enumerate(new_my):
            f.write(f"[{i}] {s['text']}\n")

    # bilingual.txt
    p = os.path.join(OUT, f"{PREFIX}_bilingual.txt"); bak(p)
    with open(p, "w", encoding="utf-8") as f:
        for i in range(len(new_my)):
            f.write(f"[{i}] {new_my[i]['text']}\n")
            f.write(f"     {new_vi[i]['text']}\n")
            f.write("\n")

    # batches
    n = (len(new_my) + BATCH_SIZE - 1) // BATCH_SIZE
    bdir = os.path.join(OUT, "batches")
    for b in range(n):
        s0 = b * BATCH_SIZE
        e0 = min(s0 + BATCH_SIZE, len(new_my))
        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(s0, e0):
                f.write(f"[{i}] {new_my[i]['text']}\n")
        with open(os.path.join(bdir, f"{name}_vi.txt"), "w", encoding="utf-8") as f:
            for i in range(s0, e0):
                f.write(f"[{i}] {new_vi[i]['text']}\n")

    print(f"✅ 001_final.srt / _vi.srt / srt_map / myanmar_lines / bilingual")
    print(f"✅ {n} batches regenerated")


if __name__ == "__main__":
    main()