#!/usr/bin/env python3
import importlib.util, json, re, shutil, time, urllib.request, subprocess, sys
from pathlib import Path

PROJECT=Path('/data/video-pipeline/Youtube-Video-Maker/Project/28052026-002-Tran-Quang-Dieu-Bui-Thi-Xuan')
SCENES=PROJECT/'script/scenes.json'
AUDIO_DIR=PROJECT/'audio/scenes'
BACKUP_DIR=PROJECT/'audio/scene_backups'
ASR_DIR=PROJECT/'subtitles/scene_pronunciation_asr_v2'
REPORT=PROJECT/'logs/step07_pronunciation_retry_v2_report.json'
STEP7=PROJECT/'scripts_run_step_07_tts_v2.py'
STEP8=PROJECT/'scripts_run_step_08_compose_tts.py'
ASR_ENDPOINT='http://192.168.1.103:8001/transcribe-and-save'
TARGETS=['scene_016','scene_047','scene_060','scene_070','scene_079','scene_093','scene_112','scene_138','scene_175','scene_193']
WATCH={'trần':['tần'],'nhưng':['nưng'],'đây':['dây']}

def log(msg): print(time.strftime('%Y-%m-%d %H:%M:%S ')+msg, flush=True)
def norm(s):
    s=s.lower(); s=re.sub(r'[^a-zà-ỹđ\s]',' ',s); return re.sub(r'\s+',' ',s).strip()
def post_json(url,payload,timeout=1800):
    req=urllib.request.Request(url,data=json.dumps(payload,ensure_ascii=False).encode(),headers={'Content-Type':'application/json'},method='POST')
    with urllib.request.urlopen(req,timeout=timeout) as r: return json.loads(r.read().decode())
def asr(scene_id):
    out=ASR_DIR/f'{scene_id}.json'; audio=AUDIO_DIR/f'{scene_id}.wav'
    payload={'input_path':str(audio),'output_path':str(out),'language':'vi','task':'transcribe','word_timestamps':True}
    resp=post_json(ASR_ENDPOINT,payload)
    doc=json.loads(out.read_text(encoding='utf-8')) if out.exists() else resp
    text=doc.get('text') or doc.get('transcript') or ' '.join(seg.get('text','') for seg in doc.get('segments',[]) or [])
    return text, str(out)
def flags(narr,tr):
    nt=norm(tr); nn=norm(narr); res=[]
    for w,bads in WATCH.items():
        src=len(re.findall(rf'\b{re.escape(w)}\b',nn))
        if not src: continue
        good=len(re.findall(rf'\b{re.escape(w)}\b',nt))
        bad=[b for b in bads if re.search(rf'\b{re.escape(b)}\b',nt)]
        if bad or good<src: res.append({'word':w,'src_count':src,'asr_good_count':good,'bad_found':bad})
    return res
def pause_payload(tts):
    # Add a comma pause before risky words unless already preceded by punctuation.
    for w, repl in [('trần','Trần'),('nhưng','Nhưng'),('đây','Đây')]:
        tts=re.sub(rf'(?<![,.;:!?])\s+{w}\b', rf', {repl}', tts, flags=re.I)
        tts=re.sub(rf'^{w}\b', repl, tts, flags=re.I)
    return tts

def main():
    BACKUP_DIR.mkdir(parents=True,exist_ok=True); ASR_DIR.mkdir(parents=True,exist_ok=True); REPORT.parent.mkdir(parents=True,exist_ok=True)
    spec=importlib.util.spec_from_file_location('step7',STEP7); mod=importlib.util.module_from_spec(spec); spec.loader.exec_module(mod)
    scenes={s['scene_id']:s for s in json.loads(SCENES.read_text(encoding='utf-8'))['scenes']}
    report={'project_id':PROJECT.name,'targets':TARGETS,'variant':'comma_pause_before_watchlist_words','items':[]}
    for sid in TARGETS:
        s=scenes[sid]; old=AUDIO_DIR/f'{sid}.wav'
        before_text,before_asr=asr(sid)
        before_flags=flags(s['narration'],before_text)
        backup=BACKUP_DIR/f'{sid}_pronunciation_v2_before_{time.strftime("%Y%m%d_%H%M%S")}.wav'
        if old.exists(): shutil.copy2(old,backup)
        base,changes=mod.normalize(s['narration'])
        item={'scene_id':sid,'narration':s['narration'],'tts_text':pause_payload(base),'normalization':changes+['watchlist_comma_pause_before_word'],'output':str(old)}
        log(f'REGEN {sid}')
        regen=mod.generate(item)
        after_text,after_asr=asr(sid) if regen.get('status')=='ok' else ('','')
        after_flags=flags(s['narration'],after_text) if after_text else [{'error':'regen_failed'}]
        # If ASR got worse, restore backup.
        restored=False
        if len(after_flags)>len(before_flags) and backup.exists():
            shutil.copy2(backup,old); restored=True; log(f'RESTORE {sid}')
        report['items'].append({'scene_id':sid,'backup':str(backup),'before_flags':before_flags,'before_asr_text':before_text,'after_flags':after_flags,'after_asr_text':after_text,'regenerate_result':regen,'restored':restored,'tts_payload':item['tts_text']})
    # rebuild manifest and Step8
    all_results=[]
    for sid,s in scenes.items():
        p=AUDIO_DIR/f'{sid}.wav'
        if p.exists(): all_results.append({'scene_id':sid,'status':'ok','output':str(p),'bytes':p.stat().st_size,'duration':mod.duration(p)})
        else: all_results.append({'scene_id':sid,'status':'failed','output':str(p),'error':'missing'})
    failed=[r for r in all_results if r['status']=='failed']
    (PROJECT/'logs/step_07_tts_v2_manifest.json').write_text(json.dumps({'project_id':PROJECT.name,'step':7,'version':'v2_pronunciation_comma_pause_verify','status':'completed' if not failed else 'partial_failed','total_scenes':len(all_results),'done':len(all_results)-len(failed),'failed':failed,'total_duration_seconds':sum((r.get('duration') or 0) for r in all_results),'results':all_results},ensure_ascii=False,indent=2),encoding='utf-8')
    report['summary']={'targets':len(TARGETS),'remaining_flagged_after':sum(1 for it in report['items'] if it['after_flags']),'restored':sum(1 for it in report['items'] if it['restored'])}
    REPORT.write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding='utf-8')
    log('RUN_STEP8')
    subprocess.check_call([sys.executable,str(STEP8)],cwd=str(PROJECT),timeout=3600)
    log('DONE '+json.dumps(report['summary'],ensure_ascii=False))
if __name__=='__main__': main()