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

PROJECT = Path('/data/video-pipeline/Youtube-Video-Maker/Project/26052026-001-Ngo-Quyen')
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')
WORKERS = 6
AUDIO_DIR = PROJECT / 'audio/scenes'
LOG = PROJECT / 'logs/step_07_tts_run.log'
MANIFEST = PROJECT / 'logs/step_07_tts_manifest.json'
PROGRESS = PROJECT / 'logs/step_07_tts_progress.jsonl'

AUDIO_DIR.mkdir(parents=True, exist_ok=True)
LOG.parent.mkdir(parents=True, exist_ok=True)

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

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

def probe(path):
    data = Path(path).read_bytes()[:16]
    if data.startswith(b'RIFF') and b'WAVE' in data[:12]:
        return 'wav'
    if data.startswith(b'ID3') or data[:2] == b'\xff\xfb':
        return 'mp3'
    return 'unknown'

def generate(scene):
    sid = scene['scene_id']
    out = AUDIO_DIR / f'{sid}.wav'
    narration = scene['narration']
    # OmniVoice TTS input is lowercased to avoid uppercase Vietnamese emphasis/stretching.
    # Keep the original narration in scenes/subtitles/manifests.
    tts_text = narration.lower()
    if out.exists() and out.stat().st_size > 1000 and probe(out) == 'wav':
        return {'scene_id': sid, 'status': 'skip_exists', 'output': str(out), 'bytes': out.stat().st_size, 'text_len': len(narration)}
    ref_text = REF_TEXT.read_text(encoding='utf-8') if REF_TEXT.exists() else ''
    fields = {
        'text': tts_text,
        'language': 'vi',
        'ref_text': ref_text,
        'speed': '1.0',
        'num_step': '16',
        'guidance_scale': '2.0',
        'effect_preset': 'broadcast',
    }
    files = {'ref_audio': REF_AUDIO}
    last = None
    for attempt in range(1, 4):
        try:
            boundary, body = multipart(fields, files)
            req = urllib.request.Request(ENDPOINT, data=body, headers={'Content-Type': 'multipart/form-data; boundary=' + boundary}, method='POST')
            with urllib.request.urlopen(req, timeout=240) as r:
                blob = r.read()
                ct = r.headers.get('Content-Type', '')
            tmp = out.with_suffix('.tmp')
            tmp.write_bytes(blob)
            fmt = probe(tmp)
            if fmt != 'wav':
                raise RuntimeError(f'unexpected audio format {fmt}, content-type={ct}, first={blob[:80]!r}')
            tmp.replace(out)
            return {'scene_id': sid, 'status': 'ok', 'attempt': attempt, 'output': str(out), 'bytes': out.stat().st_size, 'text_len': len(narration), 'content_type': ct}
        except Exception as exc:
            last = str(exc)
            log(f'ERROR {sid} attempt={attempt} {last}')
            time.sleep(4 * attempt)
    return {'scene_id': sid, 'status': 'failed', 'error': last, 'output': str(out), 'text_len': len(narration)}

def main():
    scenes = json.loads((PROJECT / 'script/scenes.json').read_text(encoding='utf-8'))['scenes']
    log(f'START step7 total={len(scenes)} workers={WORKERS}')
    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs = {ex.submit(generate, s): s for s in scenes}
        for fut in concurrent.futures.as_completed(futs):
            res = fut.result()
            results.append(res)
            with PROGRESS.open('a', encoding='utf-8') as f:
                f.write(json.dumps(res, ensure_ascii=False) + '\n')
            log(f"{res['status']} {res['scene_id']} {res.get('bytes', 0)}")
    failed = [r for r in results if r['status'] == 'failed']
    manifest = {
        'project_id': PROJECT.name,
        'step': 7,
        'status': 'completed' if not failed else 'partial_failed',
        'workers': WORKERS,
        'total_scenes': len(scenes),
        'done': len(results) - len(failed),
        'failed': failed,
        'results': sorted(results, key=lambda r: r['scene_id']),
    }
    MANIFEST.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding='utf-8')
    log('DONE ' + json.dumps({k: manifest[k] for k in ['status','total_scenes','done']}, ensure_ascii=False))

if __name__ == '__main__':
    main()