#!/usr/bin/env python3
import concurrent.futures, json, re, subprocess, sys, time
from pathlib import Path
PROJECT = Path('/data/video-pipeline/Youtube-Video-Maker/Project/03072026-001-Han-Tin')
GEN_SCRIPT = Path('/home/hermes/.hermes/skills/media/gen-img-gpt/scripts/gen_img_gpt.py')
LOG = PROJECT / 'logs/step06_images_run.log'
MANIFEST = PROJECT / 'logs/step06_images_manifest.json'
WORKERS = 5

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 ref_for(scene):
    txt=(scene.get('narration') or scene.get('text') or '').lower()
    refs=[]
    def add(name):
        p=PROJECT/f'characters/{name}_ref.png'
        if p.exists() and p.stat().st_size>1000: refs.append(str(p))
    if 'hàn tín' in txt or 'han tin' in txt:
        if any(w in txt for w in ['thuở thiếu', 'trẻ', 'chợ', 'chui qua háng', 'nghèo', 'hoài âm']): add('han-tin-tre-tuoi')
        else: add('han-tin-dai-tuong')
    if 'lưu bang' in txt: add('luu-bang')
    if 'tiêu hà' in txt: add('tieu-ha')
    if 'hạng vũ' in txt: add('hang-vu')
    if 'lữ hậu' in txt or 'trường lạc' in txt: add('lu-hau')
    if 'quân hán' in txt or 'binh sĩ' in txt: add('binh-si-han')
    if 'quân sở' in txt or 'long thư' in txt: add('quan-so')
    return refs[:2]

def generate_one(scene):
    sid=scene['scene_id']; out=PROJECT/f'images/scenes/{sid}.png'
    if out.exists() and out.stat().st_size>1000:
        return {'scene_id':sid,'status':'exists','output':str(out)}
    out.parent.mkdir(parents=True, exist_ok=True)
    prompt=scene.get('visual_prompt') or scene.get('visual_beat') or scene.get('text','')
    if scene.get('type') == 'heading':
        prompt = 'Light Vietnamese old-calligraphy title-card on warm antique paper, subtle ink wash and small seal stamp accent, airy text with breathing room. Title text: ' + scene.get('text','') + '. No modern typography, no heavy black title slam.'
    cmd=[sys.executable,str(GEN_SCRIPT),'--prompt',prompt,'--output',str(out),'--size','auto','--quality','auto','--background','auto','--image-detail','high','--output-format','png','--n','1']
    for r in ref_for(scene): cmd[2:2]=['--input-ref',r]
    try:
        p=subprocess.run(cmd,text=True,capture_output=True,timeout=900)
        if p.returncode!=0:
            return {'scene_id':sid,'status':'failed','error':(p.stderr or p.stdout)[-2000:]}
        return {'scene_id':sid,'status':'generated','output':str(out),'bytes':out.stat().st_size if out.exists() else 0}
    except Exception as e:
        return {'scene_id':sid,'status':'failed','error':repr(e)}

def generate_thumb(name, prompt):
    out=PROJECT/f'output/thumbnails/{name}'
    if out.exists() and out.stat().st_size>1000: return {'name':name,'status':'exists','output':str(out)}
    cmd=[sys.executable,str(GEN_SCRIPT),'--prompt',prompt,'--output',str(out),'--size','auto','--quality','auto','--background','auto','--image-detail','high','--output-format','png','--n','1']
    p=subprocess.run(cmd,text=True,capture_output=True,timeout=900)
    return {'name':name,'status':'generated' if p.returncode==0 else 'failed','output':str(out),'error':None if p.returncode==0 else (p.stderr or p.stdout)[-2000:]}

def main():
    LOG.write_text('',encoding='utf-8')
    ref_plan=json.loads((PROJECT/'script/character_reference_plan.json').read_text(encoding='utf-8'))
    if not ref_plan.get('validation',{}).get('required_refs_exist'):
        raise SystemExit('Step 06 blocked: required Step 05 refs missing')
    data=json.loads((PROJECT/'script/dense_scene_plan.json').read_text(encoding='utf-8'))
    scenes=data['scenes']; results=[]
    with concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs={ex.submit(generate_one,s):s['scene_id'] for s in scenes}
        for fut in concurrent.futures.as_completed(futs):
            res=fut.result(); results.append(res); log(json.dumps(res,ensure_ascii=False))
            MANIFEST.write_text(json.dumps({'project_id':PROJECT.name,'step':6,'scene_results':results},ensure_ascii=False,indent=2),encoding='utf-8')
    failed=[r for r in results if r['status']=='failed']
    if failed:
        MANIFEST.write_text(json.dumps({'project_id':PROJECT.name,'step':6,'status':'needs_retry','failed':failed,'scene_results':results},ensure_ascii=False,indent=2),encoding='utf-8')
        raise SystemExit('step06_scene_failures')
    thumb_prompt='cinematic historical biography poster for Vietnamese YouTube history series, HÀN TÍN, brilliant Chu-Han era strategist, dramatic parchment manuscript texture, smoky battlefield and imperial court tension, dark sepia bronze black palette, large Han Xin portrait, readable metallic-gold Vietnamese brush lettering title: HÀN TÍN, black ink-splash grunge title panel, no modern objects, no watermark.'
    thumbs=[generate_thumb('thumb_16x9_final.png', thumb_prompt+' 16:9 horizontal 1920x1080 composition.'), generate_thumb('thumb_9x16_final.png', thumb_prompt+' TikTok-safe reduced portrait poster 1080x1280 composition, keep face and title central.')]
    MANIFEST.write_text(json.dumps({'project_id':PROJECT.name,'step':6,'status':'completed','total_scenes':len(scenes),'scene_results':results,'thumbnails':thumbs},ensure_ascii=False,indent=2),encoding='utf-8')
    log('DONE '+json.dumps({'status':'completed','scenes':len(scenes),'thumbnails':thumbs},ensure_ascii=False))
if __name__=='__main__': main()