#!/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/002-Nha-Tan-221-207-TCN')
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, muted sepia brown black charcoal watercolor manuscript feel, Qin dynasty China 3rd century BCE, historically grounded, not full-frame glossy digital poster, not fantasy splash art, no text, no watermark.'
STYLE_HEADING='16:9 horizontal old manuscript title-card on warm antique paper, subtle ink wash, Qin dynasty map fragments, bronze seal, smoke and parchment, airy breathing room, no readable text, no watermark.'
NEG='Avoid modern objects, wrong era, Qing queue hairstyle, Ming robe, fantasy armor, samurai, European medieval armor, 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 any(x in t for x in ['doanh chính','tần thủy hoàng','hoàng đế','trường sinh','sa khâu']): keys.append('doanh-chinh-tan-thuy-hoang')
    if any(x in t for x in ['lý tư','pháp gia','tiểu triện','đốt sách','chôn nho','chuẩn hóa']): keys.append('ly-tu')
    if any(x in t for x in ['mông điềm','vạn lý trường thành','hung nô','biên ải']): keys.append('mong-diem')
    if any(x in t for x in ['triệu cao','hồ hợi','tần nhị thế','chỉ hươu','con ngựa']): keys.append('trieu-cao')
    if any(x in t for x in ['trần thắng','ngô quảng','đại trạch hương','dân phu']): keys.append('tran-thang-va-ngo-quang')
    if any(x in t for x in ['lưu bang','hạng vũ','hàm dương','tử anh']): keys.append('luu-bang-va-hang-vu')
    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}\nSection mood: {scene['narration']}. Ancient Qin seal, fragmented Warring States map, bamboo slips, black Qin banners, smoke and ash. {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: late Warring States and Qin empire, 221-207 BCE, ancient Chinese courts, armies, frontier walls, bamboo slips, maps, bronze seals, peasant uprisings and the burning of Xianyang.\nNarration beat: {scene['narration']}\nVisual direction: {scene.get('visual_beat') or scene.get('visual_prompt')}\n{ref_note}One cinematic narrative scene, historically grounded Qin-era robes and lamellar armor, black banners, parchment smoke, restrained epic 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','timed out']):
            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 'TẦN THỦY HOÀNG - ĐẾ CHẾ ĐẦU TIÊN'. Large portrait of Qin Shi Huang in black imperial robe, Warring States map torn behind him, burning Xianyang, terracotta soldiers, black bronze sepia palette, parchment texture, smoky conflict background, readable metallic-gold Vietnamese brush lettering on black ink/grunge panel, historically grounded Qin dynasty, 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('doanh-chinh-tan-thuy-hoang')
    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['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()