#!/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/22062026-001-Nha-Hau-Le-Le-So')
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 historical 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, Qing dynasty costume, Chinese opera costume, European armor, fantasy armor, samurai, guns, 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()
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'])
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):
    t=text.lower(); keys=[]
    rules=[('le-loi-le-thai-to',['lê lợi','lê thái tổ','thuận thiên']),('nguyen-trai',['nguyễn trãi','bình ngô','nhân nghĩa']),('le-lai',['lê lai']),('le-thanh-tong',['lê thánh tông','hồng đức','luật hồng đức','bia tiến sĩ','tao đàn']),('nguyen-thi-lo',['nguyễn thị lộ','lệ chi viên']),('mac-dang-dung',['mạc đăng dung','nhà mạc']),('quan-minh',['quân minh','nhà minh','minh thuộc','giao chỉ','liễu thăng','mộc thạnh','vương thông']),('quan-dan-lam-son',['lam sơn','nghĩa quân','chí linh','dân mất nước']),('nho-si-hong-duc',['nho giáo','khoa cử','văn miếu','hiền tài','nho sĩ']),('quan-chiem-thanh',['chiêm thành','trà toàn','vijaya','quảng nam'])]
    for key,words in rules:
        if any(w in t for w in words): keys.append(key)
    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):
    if scene.get('type')=='heading':
        return f"{STYLE_HEADING}\nVietnamese title card text to depict accurately if possible: {scene['narration']}\nSubtle Lê sơ / Lam Sơn / Hồng Đức motifs, Đông Kinh court seal, Bình Ngô Đại Cáo, law code, scholar steles, maps. {NEG}"
    refs=scene_refs(scene.get('narration','')+' '+scene.get('summary',''))
    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 refs else ''
    return f"{STYLE_SCENE}\nHistorical context: 15th to early 16th-century Đại Việt under Nhà Hậu Lê / Lê sơ, from Lam Sơn uprising against Ming occupation to Hồng Đức reforms, Đông Kinh royal court, Văn Miếu, law codes, maps, southern frontier.\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, court robes, scholar robes, simple armor, villages, mountains, palace, maps, steles, edicts, war and reform atmosphere.\nNegative exclusions: {NEG}"
def run_one(scene):
    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),'--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',''))[:2]: cmd += ['--input-ref',str(ref)]
    last_err=''
    for attempt in range(1, int(os.environ.get('STEP06_MAX_ATTEMPTS','6')) + 1):
        p=subprocess.run(cmd,text=True,capture_output=True,timeout=1200)
        if p.returncode==0 and out.exists() and out.stat().st_size>1000:
            return {'scene_id':scene['scene_id'],'status':'ok','output':str(out),'bytes':out.stat().st_size,'attempt':attempt}
        last_err=(p.stderr or p.stdout)[-2000:]
        if '401' in last_err or '429' in last_err or 'usage limit' in last_err.lower() or 'invalidated oauth token' in last_err.lower():
            wait=int(os.environ.get('STEP06_RATE_WAIT','150')); log(f"rate_wait {scene['scene_id']} attempt={attempt} wait={wait}s"); time.sleep(wait); continue
        break
    return {'scene_id':scene['scene_id'],'status':'failed','output':str(out),'error':last_err}
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À HẬU LÊ - TỪ LAM SƠN KHỞI NGHĨA ĐẾN THỜI HỒNG ĐỨC RỰC RỠ'. Large dramatic portraits of Lê Lợi and Lê Thánh Tông, Lam Sơn mountains, Bình Ngô Đại Cáo manuscript, Hồng Đức law code, Văn Miếu steles, Đông Kinh court, dark sepia bronze black palette, parchment texture, readable metallic-gold Vietnamese brush lettering over irregular black ink-splash title panel. Historically grounded 15th-century Đại Việt / Lê sơ, not modern, not fantasy, no watermark."
def gen_thumb(kind,out):
    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']
    for k in ['le-loi-le-thai-to','le-thanh-tong']:
        p=REFS.get(k)
        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')
    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):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); 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()