#!/usr/bin/env python3
import json
import re
import time
import urllib.request
from pathlib import Path

PROJECT = Path('/data/video-pipeline/Channel-Loi-Co-Nhan/project/13062026-099-Dung-De-Mot-Loi-Khinh-Lam-Gay-Chi-Minh')
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']:
        # Keep punctuation in display text, but split into visible tokens.
        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)

    # Force real ASR timing; do not reuse clamped/computed subtitle files.
    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))
    cue_idx = 1
    while idx < n:
        take = min(WORDS_PER_CUE, n - idx)
        # Avoid a one-word trailing cue if possible.
        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 = float(group_asr[-1]['end'])
        if end <= start:
            end = start + 0.35
        cues.append({
            'idx': cue_idx,
            'scene_id': group_script[0][0],
            'start': start,
            'end': end,
            'text': ' '.join(w for _, w in group_script),
        })
        idx += take
        cue_idx += 1

    lines = []
    for c in cues:
        lines.append(str(c['idx']))
        lines.append(f"{srt_time(c['start'])} --> {srt_time(c['end'])}")
        lines.append(c['text'])
        lines.append('')
    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()