#!/usr/bin/env python3
import base64
import json
import mimetypes
import os
import re
import sys
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from threading import Lock
from pathlib import Path

PROJECT = Path('/data/video-pipeline/Youtube-Video-Maker/Project/26052026-001-Ngo-Quyen')
ROOT = Path('/data/video-pipeline/Youtube-Video-Maker')
MODEL = 'cx/gpt-5.5-image'
BASE_URL = 'http://192.168.40.11:20128'
LOG_DIR = PROJECT / 'logs'
LOG_DIR.mkdir(parents=True, exist_ok=True)
RUN_LOG = LOG_DIR / 'step_06_run.log'
STATE_PATH = LOG_DIR / 'step_06_generation_state.json'
MAX_WORKERS = 6
LOG_LOCK = Lock()
STATE_LOCK = Lock()


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


def load_env():
    env = os.environ.copy()
    for env_file in [ROOT / '.env', PROJECT / '.env']:
        if env_file.exists():
            for line in env_file.read_text(encoding='utf-8').splitlines():
                line = line.strip()
                if not line or line.startswith('#') or '=' not in line:
                    continue
                k, v = line.split('=', 1)
                env[k.strip()] = v.strip().strip('"').strip("'")
    return env


def data_url(path):
    path = Path(path)
    mime = mimetypes.guess_type(str(path))[0] or 'image/png'
    return 'data:%s;base64,%s' % (mime, base64.b64encode(path.read_bytes()).decode())


def parse_response(text, output):
    payloads = []
    for line in text.splitlines():
        line = line.strip()
        if line.startswith('data:'):
            line = line[5:].strip()
        if not line or line == '[DONE]':
            continue
        try:
            payloads.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    if not payloads:
        payloads = [json.loads(text)]
    for payload in payloads:
        if isinstance(payload, dict) and payload.get('event') == 'error':
            raise RuntimeError(str(payload))
        data = payload.get('data') if isinstance(payload, dict) else None
        if not data:
            continue
        item = data[0] if isinstance(data, list) else data
        if not isinstance(item, dict):
            continue
        b64 = item.get('b64_json') or item.get('base64') or item.get('image_base64')
        url = item.get('url') or item.get('image_url')
        if b64:
            if b64.startswith('data:'):
                b64 = b64.split(',', 1)[1]
            Path(output).parent.mkdir(parents=True, exist_ok=True)
            Path(output).write_bytes(base64.b64decode(b64))
            return
        if url:
            Path(output).parent.mkdir(parents=True, exist_ok=True)
            urllib.request.urlretrieve(url, output)
            return
    raise RuntimeError('No image in response: ' + text[-1000:])


def generate(prompt, output, refs=None, size='1536x864', max_attempts=3):
    output = Path(output)
    if output.exists() and output.stat().st_size > 1000:
        log(f'SKIP exists {output}')
        return True
    refs = refs or []
    body = {
        'model': MODEL,
        'prompt': prompt,
        'n': 1,
        'size': size,
        'quality': 'auto',
        'background': 'auto',
        'image_detail': 'high',
        'output_format': 'png',
    }
    if refs:
        body['image'] = [data_url(r) for r in refs if Path(r).exists()]
    api_key = ENV.get('GEN_IMG_API_KEY')
    if not api_key:
        raise RuntimeError('GEN_IMG_API_KEY missing')
    url = ENV.get('GEN_IMG_BASE_URL', BASE_URL).rstrip('/') + '/v1/images/generations'
    for attempt in range(1, max_attempts + 1):
        try:
            req = urllib.request.Request(
                url,
                data=json.dumps(body).encode(),
                headers={'Content-Type': 'application/json', 'Authorization': 'Bearer ' + api_key, 'Accept': 'text/event-stream'},
                method='POST',
            )
            with urllib.request.urlopen(req, timeout=420) as response:
                text = response.read().decode('utf-8', 'ignore')
            parse_response(text, output)
            log(f'OK {output}')
            return True
        except Exception as exc:
            log(f'ERROR attempt={attempt} output={output} err={exc}')
            time.sleep(8 * attempt)
    return False


def is_heading(scene):
    return scene.get('scene_type') == 'heading' or scene.get('type') == 'heading'


def narration(scene):
    return scene.get('narration') or scene.get('text') or scene.get('summary') or ''


def scene_prompt(scene):
    text = narration(scene)
    if is_heading(scene):
        return f"""Handcrafted Vietnamese brush-calligraphy historical title card, 16:9 horizontal composition/canvas, centered on warm antique parchment paper. The visual content should occupy only about 65-70% of the page, leaving around 30-35% negative space so the page can breathe. Do not fill the whole frame. Keep the composition light, airy, elegant, handcrafted, and historically calm.
Expressive Vietnamese brush lettering with natural ink bleed, soft black-brown / charcoal ink variation, gentle wash behavior, aged paper texture, restrained splash or drip details mostly near the edges, and optional subtle seal-stamp accents. The title should not be oversized or dominate the whole page. Avoid thick fully saturated black coverage.
Show only the chapter heading text: {text}. Preserve the wording, but remove ordinal numbering if needed.
Not a modern typography poster, not a clean/plain centered text card, not an aggressive heavy-ink title slam, not a subtitle card, not a logo card, not a watermark, not edge-to-edge full-frame artwork."""
    battle_terms = ['bạch đằng', 'thuyền', 'quân', 'trận', 'cọc', 'nam hán', 'lưu hoằng tháo', 'đánh', 'giặc']
    battle = any(t in text.lower() for t in battle_terms)
    base = 'Vietnamese historical ink-wash storytelling battle illustration' if battle else 'Vietnamese historical ink-wash storytelling illustration'
    atmosphere = 'stormy historical atmosphere' if battle else 'calm historical atmosphere'
    return f"""{base}, 16:9 horizontal composition/canvas, centered irregular ink-wash window on warm antique cream paper. The painted narrative image should occupy only about 65-70% of the page, leaving around 30-35% as clean antique paper negative space for breathing room. Do not fill the entire frame edge-to-edge. Keep the page balanced, airy, elegant, and handcrafted rather than heavy or crowded. Soft brown-black ink lines, transparent watercolor wash, muted earth-tone palette, {atmosphere}.

Ground the image specifically in early 10th-century Đại Việt and the Bạch Đằng campaign world: northern Vietnamese river geography, tidal muddy banks, wooden war boats, sharpened submerged stakes when relevant, early medieval Vietnamese robe and armor language, bamboo/wood architecture, spears, bows, drums, banners, and the tense political-military atmosphere after the collapse of Tang control. Use historically appropriate Vietnamese facial features, costume language, materials, terrain, and social atmosphere. Avoid generic fantasy drift or foreign-cultural drift.

One cinematic narrative scene showing: {text}. Clear environment, clear focal action, readable emotional staging, historically grounded props and setting, visually suitable for a story video frame.

Use attached character reference only for identity consistency when applicable: face, age, hair, expression tendency, costume details, and armor language. Do not reproduce the reference sheet itself.

Not a character sheet, not a turnaround sheet, not a concept art board, not multiple views, not a lineup, not floating costume parts, not isolated armor pieces, not a design sheet, not a model sheet, not modern objects, not logo, not watermark, not poster typography, not edge-to-edge full coverage."""


def refs_for_scene(scene):
    text = narration(scene).lower()
    refs = []
    char_refs = {
        'ngô quyền': PROJECT / 'characters/ngo_quyen_ref_v3_16x9.png',
        'ngo quyen': PROJECT / 'characters/ngo_quyen_ref_v3_16x9.png',
        'kiều công tiễn': PROJECT / 'characters/kieu_cong_tien_ref_v2.png',
        'kieu cong tien': PROJECT / 'characters/kieu_cong_tien_ref_v2.png',
        'lưu hoằng tháo': PROJECT / 'characters/luu_hoang_thao_ref_v2.png',
        'luu hoang thao': PROJECT / 'characters/luu_hoang_thao_ref_v2.png',
    }
    for name, path in char_refs.items():
        if name in text and path.exists():
            refs.append(str(path))
    if not refs:
        for p in (ROOT / 'data/style-ref/main-scene').glob('*.jpg'):
            refs.append(str(p))
            if len(refs) >= 2:
                break
    return refs[:3]


def thumb_prompt(vertical=False):
    aspect = '9:16 vertical poster composition/canvas, 1080x1920' if vertical else '16:9 horizontal poster composition/canvas, 1920x1080'
    return f"""Epic dark Vietnamese historical poster, {aspect}, generated as a dedicated final thumbnail, not a crop and not a video frame. Subject: Ngô Quyền and the Battle of Bạch Đằng, series Việt Sử Kể Lại, channel Vết Mực Thời Gian. Smoky battlefield and tidal river atmosphere, hidden wooden stakes, war boats trapped in mist, storm clouds, heroic but historically grounded Ngô Quyền as central figure, sepia / bronze / gold palette, gritty parchment texture, black ink splash and grunge title panel.
Large gold Vietnamese brush-calligraphy title: VẾT MỰC THỜI GIAN. Secondary readable title: NGÔ QUYỀN. Dramatic lighting, strong face, strong contrast, clickable historical YouTube/TikTok/Facebook thumbnail style while still ancient parchment and ink-wash inspired.
Use attached Ngô Quyền character reference only for identity consistency: face, age, hair, expression tendency, costume details, armor language. Do not reproduce the reference sheet itself.
Not a character sheet, not a scene crop, not a resized scene image, not a placeholder, not modern costume, not fantasy armor, not Chinese imperial dragon robe, not watermark, not logo other than the requested title text."""


def save_state(state):
    with STATE_LOCK:
        STATE_PATH.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding='utf-8')


ENV = load_env()

def main():
    scenes_data = json.loads((PROJECT / 'script/scenes.json').read_text(encoding='utf-8'))
    scenes = scenes_data['scenes']
    existing = [s for s in scenes if (PROJECT / 'images/scenes' / f"{s.get('scene_id')}.png").exists() and (PROJECT / 'images/scenes' / f"{s.get('scene_id')}.png").stat().st_size > 1000]
    pending = [s for s in scenes if s not in existing]
    state = {
        'project_id': PROJECT.name,
        'step': 6,
        'status': 'running_batch_6',
        'workers': MAX_WORKERS,
        'total_scenes': len(scenes),
        'existing_scenes': len(existing),
        'generated_scenes': len(existing),
        'pending_scenes': len(pending),
        'failed': []
    }
    save_state(state)
    log(f'START batch scene generation workers={MAX_WORKERS} total={len(scenes)} existing={len(existing)} pending={len(pending)}')

    def run_scene(scene):
        sid = scene.get('scene_id')
        out = PROJECT / 'images/scenes' / f'{sid}.png'
        ok = generate(scene_prompt(scene), out, refs_for_scene(scene), size='1536x864')
        return sid, str(out), ok

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = [executor.submit(run_scene, scene) for scene in pending]
        for future in as_completed(futures):
            sid, out, ok = future.result()
            if ok:
                state['generated_scenes'] += 1
            else:
                state['failed'].append({'scene_id': sid, 'output': out})
            state['pending_scenes'] = max(0, state['total_scenes'] - state['generated_scenes'] - len(state['failed']))
            save_state(state)
    if not state['failed']:
        ngq = PROJECT / 'characters/ngo_quyen_ref_v3_16x9.png'
        generate(thumb_prompt(False), PROJECT / 'output/thumbnails/thumb_16x9_final.png', [str(ngq)] if ngq.exists() else [], size='1920x1080')
        generate(thumb_prompt(True), PROJECT / 'output/thumbnails/thumb_9x16_final.png', [str(ngq)] if ngq.exists() else [], size='1080x1920')
    state['status'] = 'completed' if not state['failed'] else 'completed_with_failures'
    state['thumb_16x9_exists'] = (PROJECT / 'output/thumbnails/thumb_16x9_final.png').exists()
    state['thumb_9x16_exists'] = (PROJECT / 'output/thumbnails/thumb_9x16_final.png').exists()
    save_state(state)
    log('DONE ' + json.dumps(state, ensure_ascii=False))
    return 0 if not state['failed'] else 2

if __name__ == '__main__':
    sys.exit(main())