#!/usr/bin/env python3
from __future__ import annotations

import concurrent.futures as cf
import datetime, json, os, subprocess, sys
from pathlib import Path

PROJECT=Path('/data/video-pipeline/Youtube-Video-Maker/Project/28052026-001-Tay-Son-That-Ho-Tuong')
GEN_SCRIPT=Path('/home/hermes/.hermes/skills/media/gen-img-gpt/scripts/gen_img_gpt.py')
SCENES_PATH=PROJECT/'script/scenes.json'
PLAN_PATH=PROJECT/'script/character_reference_plan.json'
LOG_PATH=PROJECT/'logs/step_06_images_batch_run.log'
STATE_PATH=PROJECT/'logs/step_06_generation_state.json'
MANIFEST_PATH=PROJECT/'logs/step_06_scene_images_manifest.json'
BATCH_SIZE=7
STYLE_SCENE='Vietnamese historical ink-wash storytelling illustration, 16:9 horizontal canvas, centered irregular ink-wash narrative window on warm antique cream paper; painted content occupies only 65-70 percent of page, remaining 30-35 percent is clean antique paper negative space; soft brown-black ink lines, transparent watercolor wash, muted earthy sepia/ochre/charcoal palette, handcrafted historical documentary mood.'
NEG='No modern objects, no watermark, no logo, no text labels unless this is a heading title card, no glossy full-frame poster, no 3D/game art, no fantasy armor, no character sheet, no concept board, no multiple views, no lineup, no model sheet.'

def log(s):
    LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
    line=datetime.datetime.now().isoformat(timespec='seconds')+' '+s
    print(line, flush=True)
    with LOG_PATH.open('a',encoding='utf-8') as f: f.write(line+'\n')

def load_env():
    ep=Path('/home/hermes/.hermes/.env')
    if ep.exists():
        for raw in ep.read_text(encoding='utf-8').splitlines():
            line=raw.strip()
            if line and not line.startswith('#') and '=' in line:
                k,v=line.split('=',1); os.environ.setdefault(k,v)

def relevant_refs(text, plan):
    tl=text.lower(); refs=[]
    for it in plan.get('items',[]):
        names=[it['name'].lower(), it['name'].split(' - ')[0].lower()]
        if any(n and n in tl for n in names): refs.append(it)
    return refs[:2]

def prompt_for_scene(scene, refs):
    if scene.get('type')=='heading':
        title=scene.get('narration') or scene.get('summary') or scene['scene_id']
        return f'Light Vietnamese old-calligraphy title-card on warm antique paper. Exact Vietnamese title text: {title}. Airy composition with breathing room, subtle ink wash, small red seal accent, elegant historical manuscript feeling. 16:9 horizontal. No extra words, no watermark, no logo, correct Vietnamese diacritics.'
    ref_txt='Use attached character reference image(s) only to preserve identity, face, age, hair, costume language and temperament. Do not reproduce the ref sheet. ' if refs else ''
    return '\n'.join([STYLE_SCENE, ref_txt+'Create one cinematic narrative scene matching this narration beat, with clear environment, action, framing, focal point, and historical grounding in late 18th century Đại Việt / Tây Sơn era.', 'Narration beat: '+scene.get('narration',''), 'Scene visual prompt: '+scene.get('visual_prompt',scene.get('visual_beat','')), 'Historical grounding: Vietnamese Tay Son period, villages, citadels, rivers, battle camps, practical Vietnamese clothing/armor/weapons appropriate to the scene.', NEG])

def gen_one(task):
    scene, plan = task
    out=PROJECT/'images/scenes'/(scene['scene_id']+'.png')
    out.parent.mkdir(parents=True,exist_ok=True)
    if out.exists() and out.stat().st_size>1000:
        return {'scene_id':scene['scene_id'],'status':'exists','path':str(out),'bytes':out.stat().st_size}
    refs=relevant_refs((scene.get('narration','')+' '+scene.get('summary','')+' '+scene.get('section','')), plan)
    ref_paths=[Path(r['output_path']) for r in refs if Path(r['output_path']).exists()]
    prompt=prompt_for_scene(scene, refs)
    cmd=[sys.executable,str(GEN_SCRIPT),'--prompt',prompt,'--output',str(out),'--size','auto','--quality','auto','--background','auto','--image-detail','high','--output-format','png']
    for r in ref_paths: cmd += ['--input-ref', str(r)]
    try:
        proc=subprocess.run(cmd,cwd=str(PROJECT),text=True,capture_output=True,timeout=480)
    except Exception as e:
        return {'scene_id':scene['scene_id'],'status':'failed','path':str(out),'error':str(e)}
    if proc.returncode==0 and out.exists() and out.stat().st_size>1000:
        return {'scene_id':scene['scene_id'],'status':'done','path':str(out),'bytes':out.stat().st_size}
    return {'scene_id':scene['scene_id'],'status':'failed','path':str(out),'error':(proc.stderr or proc.stdout or 'unknown')[-2000:]}

def run_thumbnail(name, desc, title):
    out=PROJECT/'output/thumbnails'/name; out.parent.mkdir(parents=True,exist_ok=True)
    if out.exists() and out.stat().st_size>1000: return {'thumbnail':name,'status':'exists','path':str(out)}
    prompt=f'Epic dark Vietnamese historical biography poster thumbnail, {desc}. Subject: {title}. Large central/left Tay Son warrior-general portrait group, smoky battlefield and political conflict background, dark sepia bronze black gold palette, gritty parchment texture, black ink splash/grunge title panel, large readable metallic-gold Vietnamese brush-calligraphy title: Tây Sơn Thất Hổ Tướng. GPT-generated final thumbnail, not crop, not scene image, no watermark.'
    cmd=[sys.executable,str(GEN_SCRIPT),'--prompt',prompt,'--output',str(out),'--size','auto','--quality','auto','--background','auto','--image-detail','high','--output-format','png']
    proc=subprocess.run(cmd,cwd=str(PROJECT),text=True,capture_output=True,timeout=480)
    if proc.returncode==0 and out.exists() and out.stat().st_size>1000: return {'thumbnail':name,'status':'done','path':str(out),'bytes':out.stat().st_size}
    return {'thumbnail':name,'status':'failed','path':str(out),'error':(proc.stderr or proc.stdout or 'unknown')[-2000:]}

def write_state(**kw):
    kw['updated_at']=datetime.datetime.now().isoformat(timespec='seconds')
    STATE_PATH.parent.mkdir(parents=True,exist_ok=True)
    STATE_PATH.write_text(json.dumps(kw,ensure_ascii=False,indent=2),encoding='utf-8')

def main():
    load_env()
    if not os.environ.get('GEN_IMG_API_KEY'):
        log('BLOCKED missing GEN_IMG_API_KEY'); return 2
    data=json.loads(SCENES_PATH.read_text(encoding='utf-8')); scenes=data['scenes']; title=data.get('title','Tây Sơn Thất Hổ Tướng')
    plan=json.loads(PLAN_PATH.read_text(encoding='utf-8'))
    done=[]; failed=[]
    pending=[s for s in scenes if not ((PROJECT/'images/scenes'/(s['scene_id']+'.png')).exists() and (PROJECT/'images/scenes'/(s['scene_id']+'.png')).stat().st_size>1000)]
    existing=len(scenes)-len(pending)
    log(f'START batch_size={BATCH_SIZE} existing={existing} pending={len(pending)} total={len(scenes)}')
    for start in range(0,len(pending),BATCH_SIZE):
        batch=pending[start:start+BATCH_SIZE]
        log('batch '+','.join(s['scene_id'] for s in batch))
        with cf.ThreadPoolExecutor(max_workers=BATCH_SIZE) as ex:
            futs=[ex.submit(gen_one,(s,plan)) for s in batch]
            for fut in cf.as_completed(futs):
                r=fut.result()
                if r['status'] in ('done','exists'):
                    done.append(r['scene_id']); log(f"{r['status']} {r['scene_id']} bytes={r.get('bytes','')}")
                else:
                    failed.append(r); log(f"failed {r['scene_id']}: {r.get('error','')}")
        total_done=sum(1 for s in scenes if (PROJECT/'images/scenes'/(s['scene_id']+'.png')).exists() and (PROJECT/'images/scenes'/(s['scene_id']+'.png')).stat().st_size>1000)
        write_state(step=6,phase='scenes',batch_size=BATCH_SIZE,done=total_done,total_scenes=len(scenes),failed=failed)
        if failed: break
    total_done=sum(1 for s in scenes if (PROJECT/'images/scenes'/(s['scene_id']+'.png')).exists() and (PROJECT/'images/scenes'/(s['scene_id']+'.png')).stat().st_size>1000)
    thumbs=[]
    if not failed and total_done==len(scenes):
        for name,desc in [('thumb_16x9_final.png','16:9 horizontal 1920x1080-style YouTube thumbnail'),('thumb_9x16_final.png','9:16 vertical 1080x1920-style short video thumbnail')]:
            log(f'generate thumbnail {name}')
            r=run_thumbnail(name,desc,title)
            if r['status'] in ('done','exists'):
                thumbs.append(r['path']); log(f"{r['status']} thumbnail {name}")
            else:
                failed.append(r); log(f"failed thumbnail {name}: {r.get('error','')}"); break
    status='completed' if not failed and total_done==len(scenes) and len(thumbs)==2 else 'partial_failed'
    manifest={'project_id':PROJECT.name,'step':6,'status':status,'provider':'gen-img-gpt','model':'cx/gpt-5.5-image','batch_size':BATCH_SIZE,'scene_done':total_done,'total_scenes':len(scenes),'thumbnails':thumbs,'failed':failed,'updated_at':datetime.datetime.now().isoformat(timespec='seconds')}
    MANIFEST_PATH.parent.mkdir(parents=True,exist_ok=True)
    MANIFEST_PATH.write_text(json.dumps(manifest,ensure_ascii=False,indent=2),encoding='utf-8')
    write_state(**manifest)
    log('DONE '+json.dumps({'status':status,'scene_done':total_done,'total_scenes':len(scenes),'thumbs':len(thumbs)},ensure_ascii=False))
    return 0 if status=='completed' else 1
if __name__=='__main__': raise SystemExit(main())