#!/usr/bin/env python3
import json, mimetypes, subprocess, time, urllib.request, uuid
from pathlib import Path

PROJECT = Path('/data/video-pipeline/Youtube-Video-Maker/Project/31052026-001-Nguyen-Hue-Quang-Trung')
ENDPOINT = 'http://192.168.1.201:3900/generate'
REF_AUDIO = Path('/data/video-pipeline/Youtube-Video-Maker/data/voices/template1/refaudio.mp3')
REF_TEXT = Path('/data/video-pipeline/Youtube-Video-Maker/data/voices/template1/reftext.txt')
LOG = PROJECT / 'logs/targeted_regen_7_segments.log'
REPORT = PROJECT / 'logs/targeted_regen_7_segments_report.json'
LEAD = 0.25
TAIL = 0.05

TARGETS = {
    'scene_125/seg_001': {'text': 'người chiến thắng biết nghĩ đến ngày mai.', 'speed': '0.95', 'note': 'rerun for leading artifact'},
    'scene_142/seg_001': {'text': 'năm một nghìn bảy trăm chín mươi hai, quang trung đột ngột qua đời, khi mới khoảng ba mươi chín tuổi.', 'speed': '0.92', 'note': 'spell out 1792'},
    'scene_143/seg_002': {'text': 'nguyễn ánh, sau nhiều năm bền bỉ gây dựng lực lượng ở phương nam, tận dụng thời cơ tiến lên.', 'speed': '0.92', 'note': 'protect Nguyen Anh at start'},
    'scene_149/seg_001': {'text': 'ông cùng phong trào tây sơn đánh đổ thế lực chúa nguyễn ở đàng trong, lật nhào phủ chúa trịnh ở đàng ngoài, chấm dứt cục diện chia cắt kéo dài, đánh tan quân xiêm ở rạch gầm xoài mút, đại phá quân thanh trong mùa xuân kỷ dậu.', 'speed': '1.03', 'note': 'faster, less spaced punctuation'},
    'scene_155/seg_001': {'text': 'khi ngoại bang kéo đến, ông không hề run sợ.', 'speed': '0.95', 'note': 'avoid repeated khong'},
    'scene_162/seg_001': {'text': 'có thể tưởng tượng buổi sáng vua quang trung tiến vào kinh thành.', 'speed': '1.02', 'note': 'replace colon pause, faster'},
    'scene_169/seg_001': {'text': 'tiếng vang ấy vẫn còn đó.', 'speed': '0.95', 'note': 'protect short phrase'},
}

def log(msg):
    line = time.strftime('%Y-%m-%d %H:%M:%S ') + msg
    print(line, flush=True)
    LOG.parent.mkdir(parents=True, exist_ok=True)
    with LOG.open('a', encoding='utf-8') as f:
        f.write(line + '\n')

def multipart(fields, files):
    b = '----OmniVoiceBoundary' + uuid.uuid4().hex
    chunks = []
    for k, v in fields.items():
        chunks.append(f'--{b}\r\nContent-Disposition: form-data; name="{k}"\r\n\r\n{v}\r\n'.encode())
    for k, path in files.items():
        p = Path(path)
        mime = mimetypes.guess_type(str(p))[0] or 'application/octet-stream'
        chunks.append(f'--{b}\r\nContent-Disposition: form-data; name="{k}"; filename="{p.name}"\r\nContent-Type: {mime}\r\n\r\n'.encode() + p.read_bytes() + b'\r\n')
    chunks.append(f'--{b}--\r\n'.encode())
    return b, b''.join(chunks)

def probe_fmt(path):
    data = Path(path).read_bytes()[:16]
    return data.startswith(b'RIFF') and b'WAVE' in data[:12]

def duration(path):
    try:
        out = subprocess.check_output(['ffprobe','-v','error','-show_entries','format=duration','-of','default=nw=1:nk=1',str(path)], text=True, timeout=10)
        return float(out.strip())
    except Exception:
        return None

def add_edge_silence(path):
    tmp = path.with_suffix('.padded.tmp.wav')
    subprocess.check_call(['ffmpeg','-y','-v','error','-i',str(path),'-af',f'apad=pad_dur={TAIL},adelay={int(LEAD*1000)}:all=1',str(tmp)], timeout=60)
    tmp.replace(path)

def generate(key, item, text_item):
    out = Path(item['audio_path'])
    ref_text = REF_TEXT.read_text(encoding='utf-8') if REF_TEXT.exists() else ''
    fields = {'text': text_item['text'], 'language': 'vi', 'ref_text': ref_text, 'speed': text_item['speed'], 'num_step': '16', 'guidance_scale': '2.0', 'effect_preset': 'broadcast'}
    files = {'ref_audio': REF_AUDIO}
    last = None
    for attempt in range(1, 5):
        try:
            b, body = multipart(fields, files)
            req = urllib.request.Request(ENDPOINT, data=body, headers={'Content-Type': 'multipart/form-data; boundary=' + b}, method='POST')
            with urllib.request.urlopen(req, timeout=300) as r:
                blob = r.read()
            tmp = out.with_suffix('.targeted.tmp')
            tmp.write_bytes(blob)
            if not probe_fmt(tmp):
                raise RuntimeError(f'unexpected audio format first={blob[:60]!r}')
            tmp.replace(out)
            add_edge_silence(out)
            return {'segment': key, 'status': 'ok', 'attempt': attempt, 'path': str(out), 'bytes': out.stat().st_size, 'duration': duration(out), 'payload_text': text_item['text'], 'speed': text_item['speed'], 'note': text_item['note']}
        except Exception as e:
            last = str(e)
            log(f'ERROR {key} attempt={attempt} {last}')
            time.sleep(5 * attempt)
    return {'segment': key, 'status': 'failed', 'error': last, 'path': str(out), 'payload_text': text_item['text'], 'speed': text_item['speed']}

def main():
    LOG.write_text('', encoding='utf-8')
    manifest = json.loads((PROJECT / 'audio/voice_segments/voice_segments_manifest.json').read_text(encoding='utf-8'))
    by = {f"{it['scene_id']}/seg_{it['segment_index_in_scene']:03d}": it for it in manifest['items']}
    results = []
    for key, spec in TARGETS.items():
        log(f'GENERATE {key} speed={spec["speed"]} note={spec["note"]}')
        results.append(generate(key, by[key], spec))
        log(f"RESULT {key} {results[-1]['status']} dur={results[-1].get('duration')}")
    REPORT.write_text(json.dumps({'project_id': PROJECT.name, 'status': 'completed' if all(r['status']=='ok' for r in results) else 'partial_failed', 'results': results}, ensure_ascii=False, indent=2), encoding='utf-8')
    if not all(r['status']=='ok' for r in results):
        raise SystemExit('partial_failed')

if __name__ == '__main__':
    main()