#!/usr/bin/env python3
"""
Merge per-chunk Gemini SRTs → single continuous SRT + prepare translation batches.

Project 002. Chunk offset: each chunk i starts at i * (120 - overlap) = i * 119.7s.
"""
import os, re, json
from pathlib import Path

PREFIX = "002"
PROJ = "/home/tuan-nguyen/.openclaw/workspace/015-phu_de_video"
TRANS_DIR = os.path.join(PROJ, "output", "transcripts")
OUT_DIR = os.path.join(PROJ, "output")
CHUNK_DURATION = 120.0
OVERLAP = 0.3
EFFECTIVE = CHUNK_DURATION - OVERLAP  # 119.7

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_sec(h, m, s, ms):
    return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 1000.0


def parse_srt(path):
    with open(path, encoding='utf-8') as f:
        content = f.read()
    segs = []
    for block in content.strip().split('\n\n'):
        lines = block.split('\n')
        if len(lines) < 3:
            continue
        m = TS_RE.match(lines[1].strip())
        if not m:
            continue
        start = to_sec(m.group(1), m.group(2), m.group(3), m.group(4))
        end = to_sec(m.group(5), m.group(6), m.group(7), m.group(8))
        text = '\n'.join(lines[2:]).strip()
        if text:
            segs.append((start, end, text))
    return segs


def fmt(sec):
    h = int(sec // 3600)
    m = int((sec % 3600) // 60)
    s = int(sec % 60)
    ms = int(round((sec - int(sec)) * 1000))
    if ms >= 1000:
        ms -= 1000
        s += 1
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


# ── Merge ──
all_segs = []
chunk_files = sorted(Path(TRANS_DIR).glob(f"{PREFIX}_clean_chunk_*_gemini.srt"))
print(f"Found {len(chunk_files)} chunk SRTs")

for cf in chunk_files:
    m = re.search(r'chunk_(\d+)', cf.name)
    idx = int(m.group(1))
    offset = idx * EFFECTIVE
    for (start, end, text) in parse_srt(str(cf)):
        all_segs.append((start + offset, end + offset, text))

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

# ── Clamp overlap tại ranh giới chunk (end <= start của segment kế - gap 30ms) ──
GAP_MS = 30
clamped = []
for i, (start, end, text) in enumerate(all_segs):
    if i < len(all_segs) - 1 and end > all_segs[i + 1][0] - GAP_MS:
        end = all_segs[i + 1][0] - GAP_MS
    if end <= start:
        end = start + 500
    clamped.append((start, end, text))
all_segs = clamped

# ── Write merged SRT ──
final_srt = os.path.join(OUT_DIR, f"{PREFIX}_final.srt")
with open(final_srt, 'w', encoding='utf-8') as f:
    for i, (start, end, text) in enumerate(all_segs, 1):
        f.write(f"{i}\n{fmt(start)} --> {fmt(end)}\n{text}\n\n")

# ── Extract records ──
records = []
for i, (start, end, text) in enumerate(all_segs):
    records.append({"index": i + 1, "start": fmt(start), "end": fmt(end),
                    "myanmar": text})

with open(os.path.join(OUT_DIR, f"{PREFIX}_srt_map.json"), 'w', encoding='utf-8') as f:
    json.dump(records, f, ensure_ascii=False, indent=2)

with open(os.path.join(OUT_DIR, f"{PREFIX}_myanmar_lines.txt"), 'w', encoding='utf-8') as f:
    for r in records:
        f.write(f"[{r['index'] - 1}] {r['myanmar']}\n")

# ── Batches ──
BATCH_SIZE = 30
n = (len(records) + BATCH_SIZE - 1) // BATCH_SIZE
os.makedirs(os.path.join(OUT_DIR, "batches"), exist_ok=True)
for b in range(n):
    s = b * BATCH_SIZE
    e = min(s + BATCH_SIZE, len(records))
    with open(os.path.join(OUT_DIR, "batches", f"{PREFIX}_batch_{b + 1:03d}_my.txt"),
              'w', encoding='utf-8') as f:
        for r in records[s:e]:
            f.write(f"[{r['index'] - 1}] {r['myanmar']}\n")

progress = {
    "project": PREFIX,
    "total_batches": n,
    "batches": [{"name": f"{PREFIX}_batch_{i + 1:03d}", "status": "pending"}
                for i in range(n)],
}
with open(os.path.join(OUT_DIR, "translation_progress.json"), 'w', encoding='utf-8') as f:
    json.dump(progress, f, indent=2)

print(f"{len(records)} subtitles → {n} batches")
print(f"→ {final_srt}")
print(f"→ {os.path.join(OUT_DIR, f'{PREFIX}_srt_map.json')}")