#!/usr/bin/env python3
"""Build Lời Cổ Nhân subtitles from real ASR word timestamps.

Set PROJECT to the current project path before running, or pass the project path
as argv[1]. This script intentionally refuses to copy clamped/computed subtitles.
"""
import json
import re
import sys
import time
import urllib.request
from pathlib import Path

PROJECT = Path(sys.argv[1]) if len(sys.argv) > 1 else Path('/data/video-pipeline/Channel-Loi-Co-Nhan/project/130-Giu-Goc-Re-Sach-Cay-Doi-Moi-Ben-Lau')
ENDPOINT = 'http://192.168.1.103:8001/transcribe-and-save'
HEALTH_ENDPOINT = 'http://192.168.1.103:8001/health'
MASTER_AUDIO = PROJECT / 'audio/master_narration.wav'
ASR_OUTPUT = PROJECT / 'subtitles/asr_word_timing.json'
SRT_OUTPUT = PROJECT / 'subtitles/master_subtitles_asr_timing_script_text.srt'
DEBUG_OUTPUT = PROJECT / 'subtitles/subtitle_asr_script_text_debug.json'
MANIFEST = PROJECT / 'logs/asr_subtitle_rebuild_manifest.json'
WORDS_PER_CUE = 5


def log(msg):
    print(time.strftime('%F %T ') + msg, flush=True)


def post_json(url, payload=None, timeout=7200):
    data = None if payload is None else json.dumps(payload, ensure_ascii=False).encode('utf-8')
    req = urllib.request.Request(url, data=data, headers={'Content-Type': 'application/json'} if payload is not None else {}, method='POST' if payload is not None else 'GET')
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        return json.loads(resp.read().decode('utf-8'))


def collect_words(doc):
    words = []
    for seg in doc.get('segments', []) or []:
        for w in seg.get('words', []) or []:
            if w.get('start') is not None and w.get('end') is not None:
                words.append(w)
    if not words and isinstance(doc.get('words'), list):
        words = [w for w in doc['words'] if w.get('start') is not None and w.get('end') is not None]
    return words


def srt_time(sec):
    ms = int(round(max(0.0, sec) * 1000))
    h = ms // 3600000; ms %= 3600000
    m = ms // 60000; ms %= 60000
    s = ms // 1000; ms %= 1000
    return f'{h:02}:{m:02}:{s:02},{ms:03}'


def script_words_by_scene():
    data = json.loads((PROJECT / 'script/scenes.json').read_text(encoding='utf-8'))
    items = []
    for sc in data['scenes']:
        words = re.findall(r'\S+', sc['narration'].replace('\n', ' '))
        items.extend((sc['scene_id'], w) for w in words)
    return items


def main():
    if not MASTER_AUDIO.exists():
        raise SystemExit(f'missing master audio: {MASTER_AUDIO}')
    ASR_OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    MANIFEST.parent.mkdir(parents=True, exist_ok=True)
    health = post_json(HEALTH_ENDPOINT, None, timeout=30)
    payload = {'input_path': str(MASTER_AUDIO), 'output_path': str(ASR_OUTPUT), 'language': 'vi', 'task': 'transcribe', 'word_timestamps': True}
    log(f'POST ASR {ENDPOINT}')
    response = post_json(ENDPOINT, payload, timeout=7200)
    if not ASR_OUTPUT.exists():
        ASR_OUTPUT.write_text(json.dumps(response, ensure_ascii=False, indent=2), encoding='utf-8')
    asr_doc = json.loads(ASR_OUTPUT.read_text(encoding='utf-8'))
    asr_words = collect_words(asr_doc)
    script_words = script_words_by_scene()
    if not asr_words:
        raise SystemExit('ASR returned no word timestamps')
    if len(asr_words) < len(script_words) * 0.75:
        raise SystemExit(f'ASR word count too low: asr={len(asr_words)} script={len(script_words)}')
    cues = []
    idx = 0
    n = min(len(asr_words), len(script_words))
    while idx < n:
        take = min(WORDS_PER_CUE, n - idx)
        if n - idx == WORDS_PER_CUE + 1:
            take = 3
        group_asr = asr_words[idx:idx + take]
        group_script = script_words[idx:idx + take]
        start = float(group_asr[0]['start'])
        end = max(float(group_asr[-1]['end']), start + 0.35)
        cues.append({'idx': len(cues) + 1, 'scene_id': group_script[0][0], 'start': start, 'end': end, 'text': ' '.join(w for _, w in group_script)})
        idx += take
    lines = []
    for c in cues:
        lines += [str(c['idx']), f"{srt_time(c['start'])} --> {srt_time(c['end'])}", c['text'], '']
    SRT_OUTPUT.write_text('\n'.join(lines), encoding='utf-8')
    DEBUG_OUTPUT.write_text(json.dumps({'asr_words': len(asr_words), 'script_words': len(script_words), 'ratio': len(asr_words) / max(1, len(script_words)), 'cues': cues}, ensure_ascii=False, indent=2), encoding='utf-8')
    MANIFEST.write_text(json.dumps({'status': 'completed', 'endpoint': ENDPOINT, 'health': health, 'input_audio_path': str(MASTER_AUDIO), 'asr_output': str(ASR_OUTPUT), 'srt_output': str(SRT_OUTPUT), 'asr_words': len(asr_words), 'script_words': len(script_words), 'cues': len(cues), 'note': 'Subtitle timing rebuilt from real ASR word timestamps; display text comes from original script.'}, ensure_ascii=False, indent=2), encoding='utf-8')
    log(f'DONE ASR subtitle cues={len(cues)} asr_words={len(asr_words)} script_words={len(script_words)}')


if __name__ == '__main__':
    main()