#!/usr/bin/env python3
import importlib.util, json, re, shutil, subprocess, sys, time, urllib.request
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'
TMP_DIR=PROJECT/'audio/pronunciation_candidates'
BACKUP_DIR=PROJECT/'audio/scene_backups'
ASR_DIR=PROJECT/'subtitles/scene_pronunciation_candidates_asr'
REPORT=PROJECT/'logs/step07_specific_pronunciation_candidates_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_047':'trần',
    'scene_060':'nhưng',
    'scene_070':'trần',
    'scene_093':'trần',
    'scene_112':'nhưng',
}
BAD={'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(audio_path, out):
    payload={'input_path':str(audio_path),'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
    return doc.get('text') or doc.get('transcript') or ' '.join(seg.get('text','') for seg in doc.get('segments',[]) or [])
def score(expected, tr):
    nt=norm(tr); bads=BAD[expected]
    bad_count=sum(len(re.findall(rf'\b{re.escape(b)}\b',nt)) for b in bads)
    good_count=len(re.findall(rf'\b{re.escape(expected)}\b',nt))
    return {'bad_count':bad_count,'good_count':good_count,'score':bad_count*10-good_count,'transcript':tr}
def variants(base, word):
    outs=[]
    def add(name, text): outs.append((name,text))
    if word=='trần':
        add('base_lower', base)
        add('title_case_word', re.sub(r'\btrần\b','Trần',base,flags=re.I))
        add('comma_before_title', re.sub(r'(?<![,.;:!?])\s+trần\b', ', Trần', base, flags=re.I))
        add('pause_between_name', re.sub(r'\btrần\s+quang\s+diệu\b','Trần, Quang Diệu',base,flags=re.I))
        add('period_before_title', re.sub(r'(?<![.!?])\s+trần\b', '. Trần', base, flags=re.I))
        add('linebreak_before_title', re.sub(r'\btrần\b','\nTrần',base,flags=re.I))
    elif word=='nhưng':
        add('base_lower', base)
        add('title_case_word', re.sub(r'\bnhưng\b','Nhưng',base,flags=re.I))
        add('comma_before_title', re.sub(r'(?<![,.;:!?])\s+nhưng\b', ', Nhưng', base, flags=re.I))
        add('period_before_title', re.sub(r'(?<![.!?])\s+nhưng\b', '. Nhưng', base, flags=re.I))
        add('linebreak_before_title', re.sub(r'\bnhưng\b','\nNhưng',base,flags=re.I))
        add('double_nh_soft', re.sub(r'\bnhưng\b','nhưng,',base,flags=re.I))
    return outs

def main():
    TMP_DIR.mkdir(parents=True,exist_ok=True); 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,'items':[]}
    for sid,word in TARGETS.items():
        s=scenes[sid]; dest=AUDIO_DIR/f'{sid}.wav'; backup=BACKUP_DIR/f'{sid}_manual_bad_{time.strftime("%Y%m%d_%H%M%S")}.wav'
        if dest.exists(): shutil.copy2(dest,backup)
        base,changes=mod.normalize(s['narration'])
        cands=[]
        for idx,(vname,tts) in enumerate(variants(base,word),1):
            out=TMP_DIR/f'{sid}_{idx:02d}_{vname}.wav'
            # step7.generate always writes AUDIO_DIR/<scene_id>.wav, so point it at
            # the candidate directory and use a candidate-specific scene_id.
            old_audio_dir = mod.AUDIO_DIR
            mod.AUDIO_DIR = TMP_DIR
            cand_id = out.stem
            item={'scene_id':cand_id,'narration':s['narration'],'tts_text':tts,'normalization':changes+[f'candidate_{vname}'],'output':str(out)}
            log(f'GEN {sid} {vname}')
            try:
                r=mod.generate(item)
            finally:
                mod.AUDIO_DIR = old_audio_dir
            if r.get('status')!='ok' or not out.exists():
                cands.append({'variant':vname,'result':r,'failed':True}); continue
            tr=asr(out, ASR_DIR/f'{sid}_{idx:02d}_{vname}.json')
            sc=score(word,tr)
            cands.append({'variant':vname,'result':r,'score':sc,'path':str(out),'tts_text':tts})
            log(f'ASR {sid} {vname} bad={sc["bad_count"]} good={sc["good_count"]}')
        valid=[c for c in cands if not c.get('failed')]
        best=min(valid, key=lambda c:(c['score']['score'], -c['score']['good_count'])) if valid else None
        chosen=None
        if best:
            shutil.copy2(best['path'], dest); chosen=best['variant']; log(f'CHOOSE {sid} {chosen}')
        report['items'].append({'scene_id':sid,'word':word,'backup':str(backup),'chosen':chosen,'candidates':cands})
    # rebuild manifest
    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_specific_pronunciation_candidates','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.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')
if __name__=='__main__': main()