#!/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/001-La-Van-Cau')
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','5'))
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 manuscript feeling, muted brown black charcoal olive 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 Vietnamese old-calligraphy title-card on warm antique paper, subtle ink wash, Cao Bang mountain and Dong Khe Route 4 motifs, airy breathing room, not heavy black title slam, not clean modern typography.'
NEG='Avoid modern objects, wrong culture, American Vietnam War gear, fantasy armor, samurai, Chinese opera costume, graphic gore, 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 'la văn cầu' in t or 'người lính trẻ' in t or 'tuổi thơ' in t or 'trung đoàn 174' in t: keys.append('la-van-cau-nguoi-linh-tre')
    if 'bộc phá' in t or 'cánh tay' in t or 'bị thương' in t or 'chặt' in t or 'xung phong' in t: keys.append('la-van-cau-om-boc-pha')
    if 'tổ bộc phá' in t or 'hàng rào' in t or 'lô cốt' in t or 'cửa mở' in t: keys.append('to-boc-pha-trung-doan-174')
    if 'dân công' in t or 'đồng đội' in t or 'cao bằng' in t or 'việt bắc' in t: keys.append('dong-doi-va-dan-cong-cao-bang')
    if 'quân pháp' in t or 'thực dân pháp' in t or 'địch' in t or 'cứ điểm' in t: keys.append('quan-phap-tai-dong-khe')
    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 silhouettes of Viet Bac mountains, Route 4, barbed wire, old parchment, red seal accent. {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: Vietnamese resistance against French colonial forces, Cao Bang border region, Viet Bac forests, Route 4, Battle of Dong Khe, Border Campaign Autumn-Winter 1950, Viet Minh infantry and sapper assault against fortified bunkers.\nNarration beat: {scene['narration']}\nVisual direction: {scene.get('visual_beat') or scene.get('visual_prompt')}\n{ref_note}One cinematic narrative scene, historically grounded 1950 Viet Minh uniforms, Tay mountain homeland, French colonial bunkers, barbed wire, explosive charges, smoke, dust, restrained heroic realism.\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 any(x in last_err.lower() for x in ['401','429','502','usage limit','invalidated oauth token','empty response']):
            wait=int(os.environ.get('STEP06_RATE_WAIT','120')); 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}. Dramatic Vietnamese historical biography poster for 'LA VĂN CẦU - CÁNH TAY ĐỂ LẠI BÊN CỬA MỞ ĐƯỜNG'. Large portrait of young Viet Minh hero La Văn Cầu, Cao Bằng mountains, Đông Khê bunker, barbed wire, smoke and antique parchment texture, dark sepia bronze black palette, readable metallic-gold Vietnamese brush lettering over irregular black ink-splash title panel. Historically grounded First Indochina War 1950, 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('la-van-cau-om-boc-pha') or REFS.get('la-van-cau-nguoi-linh-tre')
    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 update_check(status):
    check=json.loads((PROJECT/'checklist_19_steps.json').read_text(encoding='utf-8'))
    for s in check['steps']:
        if s['step']==6:
            s['name']='Scene Images And Thumbnails'; s['status']=status; s['completed_at']=time.strftime('%Y-%m-%dT%H:%M:%S') if status=='completed' else None
    (PROJECT/'checklist_19_steps.json').write_text(json.dumps(check,ensure_ascii=False,indent=2),encoding='utf-8')
    (PROJECT/'checklist_19_steps.md').write_text('# Checklist 19 Steps\n\n'+'\n'.join(f"- [{'x' if x['status']=='completed' else ' '}] Step {x['step']:02d}: {x['name']} - {x['status']}" for x in check['steps'])+'\n',encoding='utf-8')
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')
    update_check('completed' if status=='completed' else 'pending')
    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()