#!/usr/bin/env python3
import base64
import json
import os
import queue
import threading
import time
import urllib.error
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']
OUTDIR = PROJECT / 'images/scenes'
OUTDIR.mkdir(parents=True, exist_ok=True)
LOG = PROJECT / 'logs/step06_generate_missing_scenes.log'
LOG.parent.mkdir(parents=True, exist_ok=True)
WORKERS = 5
STOP = threading.Event()
LOCK = threading.Lock()

def log(msg):
    line = f'{time.strftime("%Y-%m-%d %H:%M:%S")} {msg}'
    with LOCK:
        print(line, flush=True)
        with LOG.open('a', encoding='utf-8') as f:
            f.write(line + '\n')

# 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'}
work = queue.Queue()
for scene in SCENES:
    out = OUTDIR / f"{scene['scene_id']}.png"
    if out.exists() and out.stat().st_size > 0:
        log(f"SKIP existing {scene['scene_id']}")
    else:
        work.put(scene)

def generate(scene):
    sid = scene['scene_id']
    out = OUTDIR / f'{sid}.png'
    prompt = scene['visual_prompt']
    body = {'model': 'cx/gpt-5.5-image', 'prompt': prompt, 'size': '1024x1792', 'n': 1}
    attempt = 0
    while not STOP.is_set():
        attempt += 1
        try:
            log(f'Generating {sid} attempt {attempt} -> {out}')
            req = urllib.request.Request(url, data=json.dumps(body).encode('utf-8'), headers=headers, method='POST')
            with urllib.request.urlopen(req, timeout=600) 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: {str(payload)[:300]}')
            if b64.startswith('data:'):
                b64 = b64.split(',', 1)[1]
            tmp = out.with_suffix('.png.tmp')
            tmp.write_bytes(base64.b64decode(b64))
            tmp.replace(out)
            log(f'DONE {sid} {out}')
            return
        except (TimeoutError, urllib.error.URLError, urllib.error.HTTPError, RuntimeError, json.JSONDecodeError) as e:
            log(f'RETRY {sid} attempt {attempt}: {type(e).__name__}: {e}')
            time.sleep(min(30 * attempt, 180))

def worker(idx):
    while not STOP.is_set():
        try:
            scene = work.get_nowait()
        except queue.Empty:
            return
        try:
            generate(scene)
        finally:
            work.task_done()

threads = [threading.Thread(target=worker, args=(i,), daemon=True) for i in range(WORKERS)]
for t in threads:
    t.start()
for t in threads:
    t.join()
log('ALL DONE')