#!/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/01062026-003-Tran-Quoc-Toan')
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.'
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 quốc toản' in t or 'hoài văn' in t or 'quả cam' in t or 'sáu chữ' in t or 'lá cờ' in t:
        keys.append('tran-quoc-toan')
    if 'cờ' in t or 'phá cường địch' in t or 'đội quân' in t or 'gia nô' in t or 'tuổi trẻ' in t:
        keys.append('six-word-flag-army')
    if 'trần nhân tông' in t or 'trần thánh tông' in t or 'thượng hoàng' in t:
        keys.append('tran-nhan-tong-thanh-tong')
    if 'trần quốc tuấn' in t or 'hưng đạo' in t or 'tiết chế' in t:
        keys.append('tran-quoc-tuan-hung-dao-vuong')
    if 'bình than' in t or 'vương hầu' in t or 'tướng lĩnh' in t or 'hội nghị' in t:
        keys.append('tran-generals-court')
    if 'nguyên' in t or 'mông' in t or 'thoát hoan' in t or 'toa đô' in t or 'ô mã nhi' in t or 'giặc' in t:
        keys.append('nguyen-mong-army')
    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 Trần dynasty wartime motifs, river Bình Than, young warrior flag silhouette, 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: 13th-century Đại Việt under the Trần dynasty, Bình Than river council, young Hoài Văn hầu Trần Quốc Toản, six-word flag, Trần army camps, anti-Nguyên resistance atmosphere.\nNarration beat: {scene['narration']}\nVisual direction: {scene.get('visual_beat') or scene.get('visual_prompt')}\n{ref_note}One cinematic narrative scene, historically grounded Vietnamese Trần dynasty clothing, armor, spears, boats, court, rivers, villages, 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)]
    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 'TRẦN QUỐC TOẢN - LÁ CỜ THÊU SÁU CHỮ'. Large dramatic bust portrait of young Hoài Văn hầu Trần Quốc Toản beside a six-word war flag, Bình Than river council and Trần army against Nguyên Mông in background, manuscript parchment texture, dark sepia bronze black palette, large readable metallic-gold Vietnamese brush lettering over irregular black ink-splash title panel. Keep hands out of frame or hidden behind cloak/flag; do not show orange fruit, squeezed fruit, close-up fingers, malformed hands, extra fingers, or gripping gesture. Historically grounded 13th-century Đại Việt Trần dynasty, 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-quoc-toan')
    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()