#!/usr/bin/env python3
import base64
import json
import os
import urllib.request
from pathlib import Path

PROJECT = Path('/data/video-pipeline/Channel-Loi-Co-Nhan/project/11062026-001-Biet-Du-Thi-Khong-Nhuc')
SCENES = json.loads((PROJECT / 'script/scenes.json').read_text(encoding='utf-8'))['scenes'][:3]
OUTDIR = PROJECT / 'images/scenes'
OUTDIR.mkdir(parents=True, exist_ok=True)

# Load GEN_IMG_API_KEY from Hermes .env without printing it.
env_path = Path('/home/hermes/.hermes/.env')
if env_path.exists():
    for line in env_path.read_text(encoding='utf-8').splitlines():
        line = line.strip()
        if not line or line.startswith('#') or '=' not in line:
            continue
        key, val = line.split('=', 1)
        if key == 'GEN_IMG_API_KEY' and not os.environ.get(key):
            os.environ[key] = val.strip().strip('"').strip("'")

api_key = os.environ.get('GEN_IMG_API_KEY')
if not api_key:
    raise SystemExit('GEN_IMG_API_KEY missing')

url = 'http://192.168.40.11:20128/v1/images/generations'
headers = {
    'Authorization': 'Bearer ' + api_key,
    'Content-Type': 'application/json',
}
for scene in SCENES:
    sid = scene['scene_id']
    out = OUTDIR / f'{sid}.png'
    prompt = scene['visual_prompt']
    print(f'Generating {sid} -> {out}', flush=True)
    body = {
        'model': 'cx/gpt-5.5-image',
        'prompt': prompt,
        'size': '1024x1792',
        'n': 1,
    }
    req = urllib.request.Request(url, data=json.dumps(body).encode('utf-8'), headers=headers, method='POST')
    with urllib.request.urlopen(req, timeout=180) as resp:
        payload = json.loads(resp.read().decode('utf-8', 'ignore'))
    item = payload['data'][0]
    b64 = item.get('b64_json') or item.get('base64') or item.get('image_base64')
    if not b64:
        raise RuntimeError(f'No b64 image for {sid}: {str(payload)[:300]}')
    if b64.startswith('data:'):
        b64 = b64.split(',', 1)[1]
    out.write_bytes(base64.b64decode(b64))
    print(str(out), flush=True)
print('DONE', flush=True)