#!/usr/bin/env python3
import concurrent.futures, json, os, subprocess, sys, time, re, unicodedata
from pathlib import Path
PROJECT=Path('/data/video-pipeline/Youtube-Video-Maker/Project/01072026-001-Nha-Tay-Son')
GEN_SCRIPT=Path('/home/hermes/.hermes/skills/media/gen-img-gpt/scripts/gen_img_gpt.py')
SCENES=PROJECT/'script/scenes.json'; PLAN=PROJECT/'script/character_reference_plan.json'
OUT_SCENES=PROJECT/'images/scenes'; OUT_THUMBS=PROJECT/'output/thumbnails'; LOG=PROJECT/'logs/step06_images_run.log'; MANIFEST=PROJECT/'logs/step06_images_manifest.json'
WORKERS=int(os.environ.get('STEP06_WORKERS','7'))
STYLE_SCENE='16:9 horizontal historical storytelling image, centered irregular ink-wash narrative window on warm antique cream paper, painted content occupies only about 65-70 percent of page with 30-35 percent clean antique paper negative space, handcrafted Vietnamese co phong feeling, muted brown black charcoal watercolor wash, not modern, not full-frame glossy digital poster, not fantasy splash art, not 3D game key art, no text, no watermark.'
STYLE_HEADING='16:9 horizontal light Vietnamese old-calligraphy title-card on warm antique paper, subtle ink wash and small red seal stamp accent, airy text with breathing room, not heavy black title slam, not clean modern typography.'
NEG='Avoid modern objects, wrong culture, fantasy armor, samurai, character sheet layout, model sheet, multiple views, floating costume parts, text labels, watermark, logo, glossy poster, edge-to-edge full-frame painting.'
def slug(s):
    s=unicodedata.normalize('NFD',s); s=''.join(c for c in s if unicodedata.category(c)!='Mn'); s=s.replace('đ','d').replace('Đ','D')
    return re.sub(r'[^a-zA-Z0-9]+','-',s).strip('-').lower()
def load_refs():
    refs={}
    if PLAN.exists():
        for r in json.loads(PLAN.read_text(encoding='utf-8')).get('references',[]): refs[r['character_id']]=Path(r['output_path'])
    return refs
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 scene_refs(text, refs):
    t=text.lower(); keys=[]
    mapping=[('nguyễn nhạc','nguyen-nhac'),('nguyễn huệ','nguyen-hue-quang-trung'),('quang trung','nguyen-hue-quang-trung'),('nguyễn lữ','nguyen-lu'),('nguyễn ánh','nguyen-anh-gia-long-tre'),('gia long','nguyen-anh-gia-long-tre'),('trương phúc loan','truong-phuc-loan'),('lê chiêu thống','le-chieu-thong'),('lê ngọc hân','cong-chua-le-ngoc-han'),('tôn sĩ nghị','quan-thanh-va-ton-si-nghi')]
    for token,k in mapping:
        if token in t: keys.append(k)
    if any(x in t for x in ['chúa nguyễn','đàng trong','phú xuân','gia định']): keys.append('chua-nguyen-cuoi-thoi')
    if any(x in t for x in ['chúa trịnh','phủ chúa','đàng ngoài','trịnh khải']): keys.append('chua-trinh-phu-chua')
    if any(x in t for x in ['vua lê','nhà lê','lê trung hưng','ngai vàng']): keys.append('vua-le-va-trieu-le-trung-hung')
    if any(x in t for x in ['quân thanh','nhà thanh','sầm nghi đống']): keys.append('quan-thanh-va-ton-si-nghi')
    if any(x in t for x in ['quân xiêm','xiêm','rạch gầm','xoài mút']): keys.append('quan-xiem')
    if any(x in t for x in ['tây sơn','áo vải','ngọc hồi','đống đa','nghĩa quân']): keys.append('quan-tay-son')
    if any(x in t for x in ['dân chúng','nông dân','dân ', 'ruộng','thuế','làng','người dân']): keys.append('dan-chung-ao-vai-tay-son')
    if any(x in t for x in ['sĩ phu','quan lại','khoa cử','nho học']): keys.append('si-phu-bac-ha-va-quan-lai')
    out=[]
    for k in keys:
        p=refs.get(k)
        if p and p.exists() and p not in out: out.append(p)
    return out
def prompt_for(scene, refs):
    if scene.get('type')=='heading':
        return f"{STYLE_HEADING}\nVietnamese title card text to depict accurately if possible: {scene['narration']}\nSubtle Tây Sơn motifs, storm banners, war drums, old citadel and river battle silhouettes, antique paper. {NEG}"
    sr=scene_refs(scene.get('narration','')+' '+scene.get('summary',''), refs)
    ref_note='Use attached character reference image(s) only to keep identity, age, face, costume language, and temperament consistent. Do not reproduce a reference sheet. ' if sr else ''
    return f"{STYLE_SCENE}\nHistorical context: late 18th century Dai Viet, Tay Son uprising and dynasty, Quy Nhon, Phu Xuan, Gia Dinh, Rach Gam - Xoai Mut, Thang Long, Ngoc Hoi - Dong Da, divided Dang Trong and Dang Ngoai.\nNarration beat: {scene['narration']}\nVisual direction: {scene.get('visual_beat') or scene.get('visual_prompt')}\n{ref_note}One cinematic narrative scene, historically grounded Vietnamese clothing, armor, banners, palaces, villages, ports, exam halls, war camps, and atmosphere.\nNegative exclusions: {NEG}"
def run_one(scene, refs):
    out=OUT_SCENES/f"{scene['scene_id']}.png"
    if out.exists() and out.stat().st_size>1000: return {'scene_id':scene['scene_id'],'status':'skipped_exists','output':str(out),'bytes':out.stat().st_size}
    cmd=[sys.executable,str(GEN_SCRIPT),'--prompt',prompt_for(scene, refs),'--output',str(out),'--size','auto','--quality','auto','--background','auto','--image-detail','high','--output-format','png','--n','1']
    for ref in scene_refs(scene.get('narration','')+' '+scene.get('summary',''), refs)[:2]: cmd += ['--input-ref',str(ref)]
    p=subprocess.run(cmd,text=True,capture_output=True,timeout=1200)
    if p.returncode!=0: return {'scene_id':scene['scene_id'],'status':'failed','output':str(out),'error':(p.stderr or p.stdout)[-2000:]}
    return {'scene_id':scene['scene_id'],'status':'ok','output':str(out),'bytes':out.stat().st_size if out.exists() else 0}
def thumb_prompt(kind):
    aspect='16:9 YouTube thumbnail, 1920x1080' if kind=='16x9' else 'TikTok-safe vertical thumbnail/poster, 1080x1280 preferred composition'
    return f"{aspect}. Cinematic Vietnamese historical dynasty poster for 'NHÀ TÂY SƠN - CƠN BÃO ÁO VẢI'. Large dramatic Quang Trung on horseback with Tây Sơn storm banners, Ngọc Hồi drums, Rạch Gầm river battle, split old order collapsing, parchment texture, dark sepia bronze black palette, large readable metallic-gold Vietnamese brush lettering over irregular black ink-splash title panel. Historically grounded, not modern, not fantasy, no watermark."
def gen_thumb(kind,out,refs):
    if out.exists() and out.stat().st_size>1000: return {'kind':kind,'status':'skipped_exists','output':str(out),'bytes':out.stat().st_size}
    cmd=[sys.executable,str(GEN_SCRIPT),'--prompt',thumb_prompt(kind),'--output',str(out),'--size','auto','--quality','auto','--background','auto','--image-detail','high','--output-format','png','--n','1']
    p=refs.get('nguyen-hue-quang-trung')
    if p and p.exists(): cmd += ['--input-ref',str(p)]
    r=subprocess.run(cmd,text=True,capture_output=True,timeout=1200)
    if r.returncode!=0: return {'kind':kind,'status':'failed','output':str(out),'error':(r.stderr or r.stdout)[-2000:]}
    return {'kind':kind,'status':'ok','output':str(out),'bytes':out.stat().st_size if out.exists() else 0}
def main():
    OUT_SCENES.mkdir(parents=True,exist_ok=True); OUT_THUMBS.mkdir(parents=True,exist_ok=True); LOG.parent.mkdir(parents=True,exist_ok=True); LOG.write_text('',encoding='utf-8')
    refs=load_refs(); scenes=json.loads(SCENES.read_text(encoding='utf-8'))['scenes']; results=[]; log(f'START step6 scenes={len(scenes)} workers={WORKERS}')
    with concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs={ex.submit(run_one,s,refs):s for s in scenes}
        for fut in concurrent.futures.as_completed(futs):
            r=fut.result(); results.append(r); log((f"{r['status']} {r.get('scene_id')} bytes={r.get('bytes')}" if 'scene_id' in r else str(r)))
    thumb_results=[]
    if not any(r['status']=='failed' for r in results):
        log('START thumbnails')
        for kind,name in [('16x9','thumb_16x9_final.png'),('9x16','thumb_9x16_final.png')]:
            tr=gen_thumb(kind,OUT_THUMBS/name,refs); thumb_results.append(tr); log(str(tr))
    status='completed' if all(r['status'] in ('ok','skipped_exists') for r in results) and len(thumb_results)==2 and all(r['status'] in ('ok','skipped_exists') for r in thumb_results) else 'needs_retry'
    MANIFEST.write_text(json.dumps({'project_id':PROJECT.name,'step':6,'status':status,'workers':WORKERS,'total_scenes':len(scenes),'scene_results':results,'thumbnail_results':thumb_results},ensure_ascii=False,indent=2),encoding='utf-8')
    log('DONE '+json.dumps({'status':status,'scenes_done':sum(1 for r in results if r['status'] in ('ok','skipped_exists')),'scenes_total':len(scenes),'thumbs':len(thumb_results)},ensure_ascii=False))
    if status!='completed': raise SystemExit('needs_retry')
if __name__=='__main__': main()