#!/usr/bin/env python3
import concurrent.futures, json, os, subprocess, sys, time
from pathlib import Path
PROJECT = Path('/data/video-pipeline/Youtube-Video-Maker/Project/05062026-003-Hai-Ba-Trung')
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, Qing/Ming costume drift, 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.'
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ưng trắc' in t or 'trưng vương' in t:
        keys.append('trung-trac-trung-vuong')
    if 'trưng nhị' in t:
        keys.append('trung-nhi')
    if 'thi sách' in t:
        keys.append('thi-sach')
    if 'tô định' in t or 'thái thú' in t:
        keys.append('to-dinh-thai-thu-giao-chi')
    if 'mã viện' in t:
        keys.append('ma-vien-tuong-dong-han')
    if 'nghĩa quân' in t or 'nữ tướng' in t or 'lê chân' in t or 'thánh thiên' in t or 'bát nàn' in t or 'thiều hoa' in t:
        keys.append('nghia-quan-me-linh-va-cac-nu-tuong')
    if 'dân' in t or 'giao chỉ' in t or 'dệt' in t or 'thuế' in t or 'ruộng' in t:
        keys.append('dan-giao-chi-duoi-ach-dong-han')
    if 'voi' in t or 'trống' in t or 'quân kỳ' in t or 'cờ' in t:
        keys.append('voi-chien-va-quan-ky-hai-ba-trung')
    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 chapter title card text to depict accurately if possible: {scene['narration']}\nSubtle Mê Linh motifs, bronze drums, elephants, Hát river, uprising banners, antique paper. {NEG}"
    text=scene.get('narration','')+' '+scene.get('summary',''); refs=scene_refs(text)
    ref_note='Use attached 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: 1st-century Giao Chỉ under Eastern Han rule, Mê Linh uprising, Lạc Việt culture, Hát river, bronze drums, elephants, Han colonial oppression and the leadership of Hai Bà Trưng.\nNarration beat: {scene['narration']}\nVisual direction: {scene.get('visual_beat') or scene.get('visual_prompt')}\n{ref_note}One cinematic narrative scene, historically grounded Lạc Việt clothing, bronze ornaments, drums, elephants, spears, banners, river villages, Mê Linh strongholds, and Eastern Han colonial contrast where needed.\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)]
    p=subprocess.run(cmd,text=True,capture_output=True,timeout=1200)
    if p.returncode!=0: return {'scene_id':scene['scene_id'],'status':'failed','output':str(out),'error':(p.stderr or p.stdout)[-2000:]}
    return {'scene_id':scene['scene_id'],'status':'ok','output':str(out),'bytes':out.stat().st_size if out.exists() else 0}
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 biography poster for 'HAI BÀ TRƯNG - TIẾNG TRỐNG MÊ LINH'. Large dramatic bust portraits of Trưng Trắc and Trưng Nhị as distinct Lạc Việt women leaders beside war elephants, bronze drums, red uprising banners, Hát river and Mê Linh atmosphere; dark sepia bronze red black palette, antique manuscript texture, large readable metallic-gold Vietnamese brush lettering over irregular black ink-splash title panel. Keep hands hidden or out of frame; no malformed fingers, no close-up gripping gesture. Historically grounded 1st-century Giao Chỉ and Lạc Việt culture, 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('trung-trac-trung-vuong')
    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')}")
    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()