#!/usr/bin/env python3
import json
import os
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

PROJECT = Path('/data/video-pipeline/Youtube-Video-Maker/Project/31052026-001-Nguyen-Hue-Quang-Trung')
GEN_SCRIPT = Path('/home/hermes/.hermes/skills/media/gen-img-gpt/scripts/gen_img_gpt.py')
LOG = PROJECT / 'logs/step06_first50_images_run.log'
MAN = PROJECT / 'logs/step06_first50_7workers_manifest.json'
LIMIT = 50
WORKERS = 7


def log(msg):
    line = f"[{datetime.now().isoformat(timespec='seconds')}] {msg}"
    print(line, flush=True)
    with LOG.open('a', encoding='utf-8') as f:
        f.write(line + '\n')


def scene_path(scene_id):
    return PROJECT / 'images/scenes' / f'{scene_id}.png'


def prompt(scene):
    n = scene.get('narration') or scene.get('text', '')
    if scene.get('type') == 'heading':
        return f"Vietnamese old-calligraphy chapter title card, 16:9 horizontal, warm antique paper, subtle ink wash, small seal stamp accent, airy breathing room. Exact title: {n}. No modern typography, no logo, no watermark."
    return f"Vietnamese historical ink-wash storytelling illustration, 16:9 horizontal canvas, centered irregular ink-wash window on warm antique cream paper, painted content only 65-70 percent of page with 30-35 percent antique paper negative space. Late 18th-century Dai Viet, Tay Son era, historically grounded Vietnamese clothing, terrain, villages, courts, armies, river geography. One cinematic narrative scene matching this narration beat: {n}. Clear focal action, environment, emotional staging. Not character sheet, not multiple views, not modern, not logo, not watermark, not text, not edge-to-edge, not glossy 3D."


def generate(scene):
    sid = scene['scene_id']
    out = scene_path(sid)
    if out.exists() and out.stat().st_size > 1000:
        return {'scene_id': sid, 'status': 'skipped_existing', 'bytes': out.stat().st_size}
    cmd = [sys.executable, str(GEN_SCRIPT), '--prompt', prompt(scene), '--output', str(out), '--size', 'auto', '--quality', 'auto', '--background', 'auto', '--image-detail', 'high', '--output-format', 'png', '--n', '1']
    proc = subprocess.run(cmd, text=True, capture_output=True, timeout=420)
    if proc.returncode != 0:
        err = (proc.stderr or proc.stdout)[-3000:]
        status = 'failed_rate_limit' if ('429' in err or 'Too Many Requests' in err or 'usage limit' in err) else 'failed'
        return {'scene_id': sid, 'status': status, 'error': err}
    if not out.exists() or out.stat().st_size < 1000:
        return {'scene_id': sid, 'status': 'failed_empty_output'}
    return {'scene_id': sid, 'status': 'generated', 'bytes': out.stat().st_size}


def write_manifest(status, results, scenes):
    remaining = [s['scene_id'] for s in scenes if not scene_path(s['scene_id']).exists()]
    data = {
        'project_id': '31052026-001-Nguyen-Hue-Quang-Trung',
        'step': '6_limited_first50_7workers_stop_on_limit',
        'workers': WORKERS,
        'limit': LIMIT,
        'status': status,
        'generated_count': LIMIT - len(remaining),
        'remaining_missing': remaining,
        'results': results,
        'note': 'User requested exactly 7 workers. Stop and report if rate limit occurs; do not switch to serial.'
    }
    MAN.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
    return data


def main():
    if not os.environ.get('GEN_IMG_API_KEY'):
        raise SystemExit('GEN_IMG_API_KEY missing')
    scenes = json.loads((PROJECT / 'script/scenes.json').read_text(encoding='utf-8'))['scenes'][:LIMIT]
    targets = [s for s in scenes if not scene_path(s['scene_id']).exists()]
    log(f'Start 7-worker missing-only run for first 50. Missing count={len(targets)}')
    if not targets:
        write_manifest('completed_limited_50_pending_review', [], scenes)
        log('Already 50/50. Stopping before thumbnails and scene_051.')
        return 0
    results = []
    status = 'completed_limited_50_pending_review'
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futures = {ex.submit(generate, s): s for s in targets}
        for fut in as_completed(futures):
            res = fut.result()
            results.append(res)
            log(f"7-worker {res['scene_id']} {res['status']}")
            if res['status'] == 'failed_rate_limit':
                status = 'stopped_rate_limit_report_to_user'
            elif res['status'].startswith('failed') and status != 'stopped_rate_limit_report_to_user':
                status = 'stopped_failure_report_to_user'
            write_manifest(status if status != 'completed_limited_50_pending_review' else 'running', results, scenes)
    final = write_manifest(status, results, scenes)
    if final['generated_count'] >= LIMIT and status == 'completed_limited_50_pending_review':
        log('Reached 50/50 images. Stopping before thumbnails and scene_051.')
        return 0
    if status == 'stopped_rate_limit_report_to_user':
        log('Rate limit encountered. Reporting to user as requested.')
        return 2
    if status == 'stopped_failure_report_to_user':
        log('Failure encountered. Reporting to user.')
        return 1
    return 0

if __name__ == '__main__':
    raise SystemExit(main())