#!/usr/bin/env python3
"""
Fix leading-silence misalignment (cách 2 — rescale chunk 000, không re-transcribe).

Vấn đề: prompt transcribe cũ ép "dòng đầu = 00:00.000" → khi video có khoảng
im lặng ở đầu, phụ đề bị neo sớm ~bằng độ dài im lặng. Cách 2 rescale tuyến tính
timestamps của chunk đầu (0..120s) sang [silence .. 120s], phần còn lại giữ nguyên.

Sửa 3 file cho mỗi project: *_final.srt, *_final_vi.srt, *_srt_map.json.
(bilingual.txt và myanmar_lines.txt không chứa timestamp → không đụng.)

Usage:
    python3 fix_leading_silence.py            # áp cho 001 & 002 theo config
    python3 fix_leading_silence.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"
OUT = os.path.join(PROJ, "output")

# silence_s = khoảng im lặng đầu đo bằng silencedetect (noise=-35dB)
# chunk0_count = số cue thuộc chunk 000 trong file merged (đếm từ chunk 000 gemini srt)
PROJECTS = {
    "001": {"silence_s": 10.8, "chunk0_count": 14},
    "002": {"silence_s": 15.9, "chunk0_count": 25},
}
CHUNK_S = 120.0  # độ dài chunk audio (giây)

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 rescale_ms(old_ms, silence_ms, chunk_ms=CHUNK_S * 1000):
    """Map [0, chunk] -> [silence, chunk] tuyến tính."""
    return silence_ms + old_ms * (chunk_ms - silence_ms) / chunk_ms


def parse_srt(path):
    """Return list of dict(start_ms, end_ms, text, raw_block)."""
    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()]
        segs.append({
            "start": to_ms(g[0], g[1], g[2], g[3]),
            "end": to_ms(g[4], g[5], g[6], g[7]),
            "text": "\n".join(lines[2:]).strip(),
        })
    return segs


def write_srt(path, segs):
    out = []
    for i, s in enumerate(segs, 1):
        out.append(f"{i}\n{fmt_ms(s['start'])} --> {fmt_ms(s['end'])}\n{s['text']}")
    with open(path, "w", encoding="utf-8") as f:
        f.write("\n\n".join(out) + "\n")


def fix_srt(path, silence_s, n0, vi=False):
    """Rescale first n0 cues; clamp boundary. vi=True → giữ gap 30ms như convention việt."""
    segs = parse_srt(path)
    assert len(segs) >= n0, f"{path}: only {len(segs)} cues < n0={n0}"
    s_ms = int(round(silence_s * 1000))
    gap_ms = 30 if vi else 0

    for i in range(n0):
        segs[i]["start"] = rescale_ms(segs[i]["start"], s_ms)
        segs[i]["end"] = rescale_ms(segs[i]["end"], s_ms)

    # boundary: cue n0-1 (chunk 000 cuối) không được tràn qua cue n0 (chunk 001 đầu)
    if n0 < len(segs):
        cap = segs[n0]["start"] - gap_ms
        if segs[n0 - 1]["end"] > cap:
            segs[n0 - 1]["end"] = cap
    if segs[n0 - 1]["end"] <= segs[n0 - 1]["start"]:
        segs[n0 - 1]["end"] = segs[n0 - 1]["start"] + 500

    write_srt(path, segs)
    return segs


def fix_srt_map(path, silence_s, n0):
    """Rescale start/end (dạng chuỗi HH:MM:SS,mmm) của n0 entry đầu."""
    with open(path, encoding="utf-8") as f:
        recs = json.load(f)
    assert len(recs) >= n0, f"{path}: only {len(recs)} entries < n0={n0}"
    s_ms = int(round(silence_s * 1000))
    ts2ms = lambda t: to_ms(*[int(x) for x in re.split(r'[:,]', t)])

    for i in range(n0):
        st = rescale_ms(ts2ms(recs[i]["start"]), s_ms)
        en = rescale_ms(ts2ms(recs[i]["end"]), s_ms)
        recs[i]["start"] = fmt_ms(st)
        recs[i]["end"] = fmt_ms(en)
    if n0 < len(recs):
        cap = ts2ms(recs[n0]["start"])
        if ts2ms(recs[n0 - 1]["end"]) > cap:
            recs[n0 - 1]["end"] = fmt_ms(cap)

    with open(path, "w", encoding="utf-8") as f:
        json.dump(recs, f, ensure_ascii=False, indent=2)
    return recs


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


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

    for prefix, cfg in PROJECTS.items():
        sil = cfg["silence_s"]
        n0 = cfg["chunk0_count"]
        files = [
            (os.path.join(OUT, f"{prefix}_final.srt"), False),
            (os.path.join(OUT, f"{prefix}_final_vi.srt"), True),
        ]
        print(f"\n=== {prefix}: silence={sil}s, chunk0={n0} cues ===")
        for path, vi in files:
            segs = parse_srt(path)
            before = segs[0]["start"] if segs else None
            if args.dry_run:
                print(f"  [dry-run] {os.path.basename(path)}: first start {fmt_ms(before)} → {fmt_ms(rescale_ms(before, int(round(sil*1000))))}")
                continue
            b = backup(path)
            fix_srt(path, sil, n0, vi=vi)
            after = parse_srt(path)[0]["start"]
            print(f"  ✅ {os.path.basename(path)}: first cue {fmt_ms(before)} → {fmt_ms(after)}  (backup: {os.path.basename(b)})")

        mp = os.path.join(OUT, f"{prefix}_srt_map.json")
        if args.dry_run:
            print(f"  [dry-run] {os.path.basename(mp)}")
        else:
            b = backup(mp)
            fix_srt_map(mp, sil, n0)
            print(f"  ✅ {os.path.basename(mp)} updated (backup: {os.path.basename(b)})")


if __name__ == "__main__":
    main()