#!/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/15062026-001-Nha-Tran')
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 'trần hưng đạo' in t or 'trần quốc tuấn' in t or 'hịch tướng sĩ' in t or 'bạch đằng' in t: keys.append('tran-hung-dao-tran-quoc-tuan')
    if 'trần nhân tông' in t or 'trúc lâm' in t or 'yên tử' in t: keys.append('tran-nhan-tong')
    if 'trần thái tông' in t or 'trần cảnh' in t or 'yên tử' in t: keys.append('tran-thai-tong-tran-canh')
    if 'trần thủ độ' in t or 'lý chiêu hoàng' in t or 'chuyển giao' in t: keys.append('tran-thu-do')
    if 'trần thánh tông' in t or 'thái thượng hoàng' in t: keys.append('tran-thanh-tong')
    if 'trần quốc toản' in t or 'bóp nát quả cam' in t or 'sáu chữ' in t: keys.append('tran-quoc-toan')
    if 'trần quang khải' in t or 'trần nhật duật' in t or 'hàm tử' in t or 'chương dương' in t or 'tây kết' in t: keys.append('tran-quang-khai-va-tran-nhat-duat')
    if 'trần khánh dư' in t or 'vân đồn' in t or 'trương văn hổ' in t: keys.append('tran-khanh-du')
    if 'nguyên mông' in t or 'nhà nguyên' in t or 'thoát hoan' in t or 'ô mã nhi' in t or 'hốt tất liệt' in t or 'toa đô' in t or 'ngột lương hợp thai' in t: keys.append('quan-nguyen-mong-thoat-hoan-o-ma-nhi')
    if 'diên hồng' in t or 'quân dân' in t or 'đại việt' in t or 'vườn không nhà trống' in t or 'bô lão' in t: keys.append('quan-dan-dai-viet-thoi-tran')
    if 'chu văn an' in t or 'thất trảm sớ' in t: keys.append('chu-van-an')
    if 'hồ quý ly' in t or 'tây đô' in t or 'nhà hồ' in t: keys.append('ho-quy-ly')
    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 Nhà Trần / Thăng Long / Bạch Đằng motifs, Đông A spirit, river stakes and bronze drums, 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: 13th-14th century Đại Việt under Nhà Trần, Thăng Long royal court, Thiên Trường, Vạn Kiếp, Diên Hồng, three Yuan-Mongol invasions, Bạch Đằng river warfare, Trúc Lâm Yên Tử, and late-Trần decline.\nNarration beat: {scene['narration']}\nVisual direction: {scene.get('visual_beat') or scene.get('visual_prompt')}\n{ref_note}One cinematic narrative scene, historically grounded 13th-century Đại Việt clothing, Trần dynasty armor, river boats, wooden stakes, Thăng Long court, village and battlefield 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,6):
        p=subprocess.run(cmd,text=True,capture_output=True,timeout=1200)
        if p.returncode==0:
            return {'scene_id':scene['scene_id'],'status':'ok','output':str(out),'bytes':out.stat().st_size if out.exists() else 0,'attempt':attempt}
        last_err=(p.stderr or p.stdout)[-2000:]
        if '429' in last_err or 'usage limit' in last_err.lower() or 'Too Many Requests' in last_err:
            time.sleep(20*attempt)
            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À TRẦN - HÀO KHÍ ĐÔNG A VÀ BA LẦN ĐÁNH BẠI NGUYÊN MÔNG'. Large dramatic portrait of Trần Hưng Đạo as a 13th-century Đại Việt commander, Bạch Đằng river stakes, Đông A banners, Mongol-Yuan invasion fleet breaking in storm clouds, manuscript parchment texture, dark sepia bronze black palette, large readable metallic-gold Vietnamese brush lettering over irregular black ink-splash title panel. Historically grounded Nhà Trần / Đại Việt, 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('tran-hung-dao-tran-quoc-tuan')
    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()