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

PROJECT=Path('/data/video-pipeline/Youtube-Video-Maker/Project/28052026-001-Tay-Son-That-Ho-Tuong')
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=2
LEADING_SILENCE_SEC=0.35
TRAILING_SILENCE_SEC=0.10
AUDIO_DIR=PROJECT/'audio/scenes'
BACKUP_DIR=PROJECT/'audio/scenes_v1_backup_before_pause_fix'
LOG=PROJECT/'logs/step_07_tts_v2_run.log'
MANIFEST=PROJECT/'logs/step_07_tts_v2_manifest.json'
TEXT_MANIFEST=PROJECT/'script/tts_text_manifest.json'
INSTRUCT=('Đọc giọng kể chuyện lịch sử Việt Nam, trầm, chậm rãi, rõ chữ. '
          'Ngắt nghỉ tự nhiên sau dấu phẩy, dấu chấm phẩy, dấu hai chấm và dấu chấm. '
          'Dừng lâu hơn giữa các câu. Không đọc liền mạch, không nuốt chữ.')
ROMAN={'I':1,'II':2,'III':3,'IV':4,'V':5,'VI':6,'VII':7,'VIII':8,'IX':9,'X':10,'XI':11,'XII':12,'XIII':13,'XIV':14,'XV':15,'XVI':16,'XVII':17,'XVIII':18,'XIX':19,'XX':20}

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

def normalize(text):
    changes=[]
    def repl(m):
        numeral=m.group(2).upper()
        val=ROMAN.get(numeral)
        if val:
            changes.append(f'{m.group(0)}->{m.group(1)} {val}')
            return f'{m.group(1)} {val}'
        return m.group(0)
    text2=re.sub(r'(?i)\b(thế\s*k[ỷỉ])\s+([IVXLCDM]{1,6})\b', repl, text)
    # OmniVoice handles colon, semicolon, and dash poorly; use comma pauses
    # only in the TTS payload while preserving scenes.json unchanged.
    before=text2
    text2=re.sub(r'\s*[:;\-–—]\s*', ', ', text2)
    text2=re.sub(r'([.!?])\s+', r'\1\n', text2)
    text2=re.sub(r',\s+', ',  ', text2)
    text2=re.sub(r'\n{2,}', '\n', text2).strip()
    if text2!=before: changes.append('punctuation_to_comma_pause')
    # OmniVoice can over-emphasize or stretch uppercase Vietnamese syllables.
    # Lowercase only the TTS payload; source text/subtitles keep original casing.
    lower = text2.lower()
    if lower != text2:
        changes.append('lowercase_for_omnivoice')
    return lower, changes

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 'wav' if data.startswith(b'RIFF') and b'WAVE' in data[:12] else 'unknown'

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):
    if LEADING_SILENCE_SEC <= 0 and TRAILING_SILENCE_SEC <= 0:
        return
    tmp = path.with_suffix('.padded.tmp.wav')
    subprocess.check_call([
        'ffmpeg', '-y', '-v', 'error', '-i', str(path),
        '-af', f'apad=pad_dur={TRAILING_SILENCE_SEC},adelay={int(LEADING_SILENCE_SEC * 1000)}:all=1',
        str(tmp),
    ], timeout=60)
    tmp.replace(path)

def generate(item):
    sid=item['scene_id']; out=AUDIO_DIR/f'{sid}.wav'
    ref_text=REF_TEXT.read_text(encoding='utf-8') if REF_TEXT.exists() else ''
    # The local TTS server only accepts a fixed taxonomy for `instruct`.
    # Keep the Vietnamese delivery guidance in the manifest, but do not send it.
    fields={'text':item['tts_text'],'language':'vi','ref_text':ref_text,'speed':'0.88','num_step':'16','guidance_scale':'2.0','effect_preset':'broadcast'}
    files={'ref_audio':REF_AUDIO}
    last=None
    for attempt in range(1,4):
        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(); ct=r.headers.get('Content-Type','')
            tmp=out.with_suffix('.tmp'); tmp.write_bytes(blob)
            if probe_fmt(tmp)!='wav': raise RuntimeError(f'unexpected format ct={ct} first={blob[:60]!r}')
            tmp.replace(out)
            add_edge_silence(out)
            return {'scene_id':sid,'status':'ok','attempt':attempt,'output':str(out),'bytes':out.stat().st_size,'duration':duration(out),'leading_silence_sec':LEADING_SILENCE_SEC,'trailing_silence_sec':TRAILING_SILENCE_SEC,'normalization':item['normalization']}
        except Exception as e:
            last=str(e); log(f'ERROR {sid} attempt={attempt} {last}'); time.sleep(5*attempt)
    return {'scene_id':sid,'status':'failed','error':last,'output':str(out)}

def main():
    scenes=json.loads((PROJECT/'script/scenes.json').read_text(encoding='utf-8'))['scenes']
    AUDIO_DIR.mkdir(parents=True, exist_ok=True)
    if not BACKUP_DIR.exists() and any(AUDIO_DIR.glob('scene_*.wav')):
        shutil.copytree(AUDIO_DIR, BACKUP_DIR)
        log(f'BACKUP {BACKUP_DIR}')
    items=[]
    for s in scenes:
        tts, changes=normalize(s['narration'])
        items.append({'scene_id':s['scene_id'],'narration':s['narration'],'tts_text':tts,'normalization':changes,'output':str(AUDIO_DIR/f"{s['scene_id']}.wav")})
    TEXT_MANIFEST.write_text(json.dumps({'project_id':PROJECT.name,'rules':['roman_century_to_number','pause_punctuation','edge_silence_padding'],'instruct':INSTRUCT,'speed':0.88,'leading_silence_sec':LEADING_SILENCE_SEC,'trailing_silence_sec':TRAILING_SILENCE_SEC,'items':items},ensure_ascii=False,indent=2),encoding='utf-8')
    log(f'START v2 total={len(items)} workers={WORKERS}')
    results=[]
    with concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs=[ex.submit(generate,it) for it in items]
        for fut in concurrent.futures.as_completed(futs):
            r=fut.result(); results.append(r); log(f"{r['status']} {r['scene_id']} {r.get('bytes',0)} dur={r.get('duration')}")
    failed=[r for r in results if r['status']=='failed']
    manifest={'project_id':PROJECT.name,'step':7,'version':'v2_pause_roman_fix','status':'completed' if not failed else 'partial_failed','total_scenes':len(items),'done':len(items)-len(failed),'failed':failed,'total_duration_seconds':sum((r.get('duration') or 0) for r in results),'results':sorted(results,key=lambda x:x['scene_id'])}
    MANIFEST.write_text(json.dumps(manifest,ensure_ascii=False,indent=2),encoding='utf-8')
    log('DONE '+json.dumps({'status':manifest['status'],'done':manifest['done'],'duration_min':round(manifest['total_duration_seconds']/60,2)},ensure_ascii=False))
if __name__=='__main__': main()