#!/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/25012026-001-Thanh-Cat-Tu-Han')
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, 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=[]
    if 'hồ quý ly' in t or 'quý ly' in t: keys.append('ho-quy-ly')
    if 'hồ hán thương' in t or 'hán thương' in t: keys.append('ho-han-thuong')
    if 'trần thiếu đế' in t or 'tôn thất nhà trần' in t or 'nhà trần' in t: keys.append('tran-thieu-de-va-ton-that-nha-tran')
    if 'chu văn an' in t or 'thất trảm sớ' in t: keys.append('chu-van-an')
    if 'đại ngu' in t or 'dân chúng' in t or 'lòng dân' in t or 'thành tây đô' in t: keys.append('quan-dan-dai-ngu')
    if 'quân minh' in t or 'nhà minh' in t or 'phù trần diệt hồ' in t: keys.append('quan-minh')
    if 'chế bồng nga' in t or 'chiêm thành' in t: keys.append('che-bong-nga-va-quan-chiem-thanh')
    if 'lê lợi' in t or 'lam sơn' in t: keys.append('le-loi-va-nghia-quan-lam-son')
    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 Mongol steppe motifs, felt tents, horse silhouettes, banners, dust storms and Eternal Blue Sky atmosphere, antique paper. {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: 12th-13th century Mongol steppe and Mongol Empire expansion, Borjigin clan, tribal rivalries, kurultai, horse archers, Jin/Xi Xia fronts and Khwarezm campaign.\nNarration beat: {scene['narration']}\nVisual direction: {scene.get('visual_beat') or scene.get('visual_prompt')}\n{ref_note}One cinematic narrative scene, historically grounded Mongol and Inner Asian clothing, felt tents, horses, composite bows, lamellar/leather armor, vast steppe, desert routes, fortified cities when relevant, conquest turmoil and war clouds.\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:]
        # Image backend sometimes returns short-lived 401/429 during token/quota rotation.
        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 world history biography poster for 'THÀNH CÁT TƯ HÃN - TỪ ĐỨA TRẺ BỊ BỎ RƠI ĐẾN CƠN BÃO LÀM RUNG CHUYỂN THẾ GIỚI'. Large dramatic portrait of Temüjin / Genghis Khan as a 13th-century Mongol khan, vast steppe, felt tents, horse archers, banners, dust storm, burning distant city silhouettes, manuscript parchment texture, dark sepia bronze black palette, large readable metallic-gold Vietnamese brush lettering over irregular black ink-splash title panel. Historically grounded 12th-13th century Mongol Empire, 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']
    p=REFS.get('temujin-thanh-cat-tu-han')
    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()