#!/usr/bin/env python3
"""
generate_previz.py — Previz Keyframe Generator (GATE 2)

Generates one hero keyframe per shot at low-res previz settings (512x896, 8 steps,
z_image_turbo). Outputs PNG frames, a previz_manifest.json, and an HTML contact-sheet
review page with lightbox navigation.

This is the cheap visual review before committing to full production generation.

Usage:
    python3 generate_previz.py PROJECT/ --episode N
    python3 generate_previz.py PROJECT/ --episode N --shots 1-5
    python3 generate_previz.py PROJECT/ --episode N --skip-existing
    python3 generate_previz.py PROJECT/ --episode N --changed-only
    python3 generate_previz.py PROJECT/ --episode N --dry-run
    python3 generate_previz.py PROJECT/ --episode N --seed 42
    python3 generate_previz.py PROJECT/ --episode N --debug-prompts
    python3 generate_previz.py PROJECT/ --episode N --no-html

Env vars:
    FAL_KEY — fal.ai API key (required)

Dependencies:
    pip install fal-client requests
"""

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional, Set

# ── Path setup ───────────────────────────────────────────────────────────

_tools_dir = str(Path(__file__).resolve().parent)
if _tools_dir not in sys.path:
    sys.path.insert(0, _tools_dir)

_lib_dir = str(Path(__file__).resolve().parent.parent / "lib")
if _lib_dir not in sys.path:
    sys.path.insert(0, _lib_dir)

# ── Imports from engine ──────────────────────────────────────────────────

from generate_storyboard_keyframes import (
    detect_characters,
    determine_generation_type,
)
from prompt_compiler import (
    OverrideStore,
    PreviousShotContext,
    _build_spatial_directive,
    _load_breakdown,
    _load_project_config,
)
from train_lora import load_registry, get_inference_config
from asset_naming import get_asset_code, build_asset_name, char_tag_from_list
from cost_tracker import CostTracker

# Conditional import — storyboard_version may not have versions yet
try:
    from storyboard_version import get_changed_shots
except ImportError:
    get_changed_shots = None

# Conditional import — visual_gate for QC loop
try:
    from visual_gate import run_gates, resolve_ref_paths, get_gemini_model
    HAS_VISUAL_GATE = True
except ImportError:
    HAS_VISUAL_GATE = False

# ── QC Constants ────────────────────────────────────────────────────────

QC_MAX_RETRIES = 2  # Max regen attempts after auto_reject

# ── Previz Constants ─────────────────────────────────────────────────────

PREVIZ_WIDTH = 512
PREVIZ_HEIGHT = 896
PREVIZ_STEPS = 8
PREVIZ_ENDPOINT = "fal-ai/z-image/turbo/lora"
PREVIZ_MODEL_KEY = "z_image"
PREVIZ_LORA_KEY = "z_image_t2i_path"
PREVIZ_LORA_SCALE_OVERRIDE = None  # Use registry scale_solo (1.3 for z-image)
PREVIZ_GUIDANCE_SCALE = None  # Z-Image Turbo doesn't use guidance
PREVIZ_MAX_WORDS = 80  # Rich prompting — full descriptions + lens language

# Character class words — used to replace physical descriptions in prompts
# so the LoRA controls identity, not the text.
CHARACTER_CLASS_WORDS = {
    "jinx": "woman",
    "kian": "armored android",
    "varek": "man",
}

# Adjectives that describe physical appearance (fight with LoRA identity)
_APPEARANCE_ADJS = (
    r"(?:young|old|lean|wiry|thin|slim|tall|short|massive|broad|huge|muscular|"
    r"beautiful|attractive|pretty|handsome|rugged|scarred|weathered|"
    r"dark-haired|red-haired|blonde|brunette|"
    r"dark-skinned|light-skinned|pale|"
    r"[a-z]+-looking|[a-z]+-featured)"
)


def _strip_character_descriptions(text: str, class_word: str) -> str:
    """Strip physical appearance descriptions, replace with class word.

    Keeps action, environment, lighting, camera — strips adjectives that
    describe what the character looks like (age, build, hair, skin) because
    the LoRA should control identity, not the text prompt.

    Examples:
        'a lean wiry woman in patched salvager gear' → 'a woman'
        'a young woman's face lit by blue light' → 'a woman's face lit by blue light'
        'a massive military combat chassis' → 'an armored android'
    """
    import re

    # Pattern: "a/the [appearance adjectives]+ [character noun phrase]"
    # Replace with just "a {class_word}"
    char_nouns = r"(?:woman|man|girl|boy|figure|person|face|chassis|android|body)"

    # Strip chains of appearance adjectives before character nouns
    text = re.sub(
        rf"\b(a|an|the|this)\s+(?:{_APPEARANCE_ADJS}\s+)*"
        rf"(?:{_APPEARANCE_ADJS}\s+)*"
        rf"({char_nouns})",
        lambda m: f"a {class_word}" if m.group(2) != "face" else f"a {class_word}'s face",
        text,
        flags=re.IGNORECASE,
    )

    # Also strip standalone appearance descriptions after character refs
    # e.g., "woman with short dark hair" → "woman"
    text = re.sub(
        rf"\b({re.escape(class_word)})\s+with\s+[\w\s]{{2,20}}(?:hair|eyes|skin|face|features|build)",
        rf"\1",
        text,
        flags=re.IGNORECASE,
    )

    # Clean double spaces
    text = re.sub(r"  +", " ", text).strip()

    return text


# ── Perspective Injection ──────────────────────────────────────────────

# Strong perspective directives keyed by camera_angle.
# These get prepended when first_frame text lacks explicit perspective cues.
# Goal: break the "subject facing camera" default. The world feels bigger
# when we see characters from behind, in profile, from extreme angles.

_PERSPECTIVE_CUES = {
    # Words in first_frame that already indicate non-frontal perspective
    "from behind", "from the back", "rear view", "over the shoulder",
    "over-the-shoulder", "three-quarter", "side profile", "silhouette",
    "seen from behind", "from below", "from above", "bird's eye",
    "POV", "point of view", "dutch angle",
}

# Shot ID mod cycle — deterministic variety so consecutive shots don't
# all get the same treatment. Uses shot ID to pick a perspective variant.
_ANGLE_VARIANTS = {
    "eye": [
        "",  # some eye-level shots stay as-is
        "Three-quarter rear view, subject angled away from camera. ",
        "Over-the-shoulder framing, camera behind the subject. ",
        "Side profile view, subject facing screen-left. ",
    ],
    "low": [
        "Extreme low angle looking up, camera near floor level. ",
        "Low angle from behind, looking up past the subject. ",
        "Low three-quarter angle, subject's back partially to camera. ",
        "Worm's eye view looking up. ",
    ],
    "high": [
        "High angle looking down, subject seen from above and behind. ",
        "Bird's eye view looking straight down. ",
        "High three-quarter angle, camera above and to the side. ",
        "Overhead angle, subject partially obscured by environment. ",
    ],
    "dutch": [
        "Dutch angle, frame tilted 15 degrees. ",
        "Canted frame, off-kilter perspective. ",
    ],
}


def _inject_perspective(first_frame: str, shot: dict,
                        prev_context: "PreviousShotContext" = None) -> str:
    """Prepend perspective language if first_frame lacks it.

    Routes through _build_spatial_directive() to get transition language,
    spatial_note always-inject, blocking directives, and camera angle
    promotion — the same system used by production pipeline.

    Falls back to _ANGLE_VARIANTS deterministic pick when the storyboard
    has no spatial/blocking/edge data.
    """
    # Check if first_frame already has strong perspective cues
    ff_lower = first_frame.lower()
    for cue in _PERSPECTIVE_CUES:
        if cue in ff_lower:
            return first_frame  # already has perspective — don't double up

    # Route through the unified spatial directive system
    directive = _build_spatial_directive(shot, model="z_image",
                                         prev_context=prev_context)
    if directive:
        return directive + " " + first_frame

    # Fallback: deterministic pick based on shot ID when no spatial data
    angle = shot.get("camera_angle", "eye").lower()
    variants = _ANGLE_VARIANTS.get(angle, _ANGLE_VARIANTS["eye"])

    shot_id = shot.get("id", 0)
    variant = variants[shot_id % len(variants)]

    if variant:
        return variant + first_frame
    return first_frame


# ── Previz Prompt Building ─────────────────────────────────────────────

def build_previz_prompt(
    shot: dict,
    characters: List[str],
    lora_registry: Dict[str, dict],
    prev_context: "PreviousShotContext" = None,
) -> str:
    """Build a short previz prompt from first_frame + LoRA trigger.

    Uses rich prompts (physical descriptions + lens language kept intact).
    Injects perspective language to break the presentational front-facing
    default — forces three-quarter, rear, over-shoulder, and extreme angles.
    Spatial-aware: routes through _build_spatial_directive for consistent
    facing, transition language, and screen direction.

    Prepends LoRA trigger word for character identity.
    Truncates to PREVIZ_MAX_WORDS to stay in attention window.
    """
    # Get trigger word and class word for primary character
    trigger = ""
    class_word = "person"
    chars_with_lora = [c for c in characters
                       if c in lora_registry and lora_registry[c].get(PREVIZ_LORA_KEY)]
    if chars_with_lora:
        primary = chars_with_lora[0]
        trigger = lora_registry[primary].get("trigger", "")
        class_word = CHARACTER_CLASS_WORDS.get(primary, "person")

    # Use first_frame as the main prompt body
    first_frame = shot.get("first_frame", "")

    if not first_frame:
        first_frame = shot.get("subject", "") or shot.get("action", "")

    # Inject perspective variety — spatial-aware with cross-shot context
    first_frame = _inject_perspective(first_frame, shot,
                                      prev_context=prev_context)

    # Build prompt: trigger first, then first_frame
    if trigger:
        prompt = f"{trigger}, {first_frame}"
    else:
        prompt = first_frame

    # Truncate to stay within attention window
    words = prompt.split()
    if len(words) > PREVIZ_MAX_WORDS:
        prompt = " ".join(words[:PREVIZ_MAX_WORDS])

    return prompt


# ── Shot Filter Parsing ──────────────────────────────────────────────────

def _parse_shots_arg(filter_str: str) -> Set[int]:
    """Parse --shots argument into a set of shot IDs.

    Supports:
        '1-5'       → {1, 2, 3, 4, 5}
        '1,3,8'     → {1, 3, 8}
        '1-3,7,9'   → {1, 2, 3, 7, 9}
        '5'         → {5}
    """
    ids: Set[int] = set()
    for part in filter_str.split(","):
        part = part.strip()
        if "-" in part:
            start, end = part.split("-", 1)
            ids.update(range(int(start), int(end) + 1))
        else:
            ids.add(int(part))
    return ids


# ── fal.ai Previz Generation ────────────────────────────────────────────

def generate_previz_frame(
    prompt: str,
    characters: List[str],
    lora_registry: Dict[str, dict],
    width: int = PREVIZ_WIDTH,
    height: int = PREVIZ_HEIGHT,
    steps: int = PREVIZ_STEPS,
    seed: int = 42,
) -> dict:
    """Generate a single previz frame via fal.ai z-image/turbo/lora.

    T2I only — no img2img, no guidance, no negative prompt (z_image turbo).

    Returns:
        fal.ai result dict with 'images' key.
    """
    import fal_client

    # Build LoRA list — single LoRA only (multiple LoRAs cause artifacts)
    loras = []
    chars_with_lora = [c for c in characters
                       if c in lora_registry and lora_registry[c].get(PREVIZ_LORA_KEY)]
    if chars_with_lora:
        # Use first character's LoRA only (primary/foreground character)
        char = chars_with_lora[0]
        reg = lora_registry[char]
        lora_path = reg[PREVIZ_LORA_KEY]
        scale = PREVIZ_LORA_SCALE_OVERRIDE if PREVIZ_LORA_SCALE_OVERRIDE else reg["scale_solo"]
        loras.append({"path": lora_path, "scale": scale})

    args = {
        "prompt": prompt,
        "image_size": {"width": width, "height": height},
        "num_inference_steps": steps,
        "seed": seed,
        "num_images": 1,
        "output_format": "png",
        "enable_safety_checker": False,
    }

    if PREVIZ_GUIDANCE_SCALE:
        args["guidance_scale"] = PREVIZ_GUIDANCE_SCALE

    if loras:
        args["loras"] = loras

    result = fal_client.subscribe(PREVIZ_ENDPOINT, arguments=args)
    return result


def download_file(url: str, output_path: str) -> str:
    """Download a file from URL to local path."""
    import requests

    resp = requests.get(url, timeout=120)
    resp.raise_for_status()
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "wb") as f:
        f.write(resp.content)
    return output_path


# ── Per-Shot Processing ──────────────────────────────────────────────────

def process_shot(
    shot: dict,
    storyboard: dict,
    output_dir: Path,
    episode: int,
    asset_code: str,
    lora_registry: Dict[str, dict],
    breakdown: dict,
    project_config: dict,
    override_store: Optional[OverrideStore],
    seed: int = 42,
    debug_prompts: bool = False,
    tracker: Optional[CostTracker] = None,
    prev_context: "PreviousShotContext" = None,
) -> dict:
    """Process a single shot: compile prompt → generate → download.

    Returns:
        Result dict with keys: shot_id, status, path, prompt_hash, characters,
        generation_type, prompt (if debug), error (if failed).
    """
    shot_id = shot["id"]
    characters = detect_characters(shot, storyboard)
    gen_type = determine_generation_type(shot)
    char_tag = char_tag_from_list(characters)
    asset_base = build_asset_name(asset_code, episode, shot_id, 1, char_tag)
    output_path = output_dir / f"{asset_base}_previz.png"

    result = {
        "shot_id": shot_id,
        "characters": characters,
        "generation_type": gen_type,
        "char_tag": char_tag,
        "asset_name": asset_base,
    }

    # Check LoRA availability — skip if no z_image LoRA for any character
    if characters:
        has_lora = any(
            c in lora_registry and lora_registry[c].get(PREVIZ_LORA_KEY)
            for c in characters
        )
        if not has_lora:
            result["status"] = "blocked_lora"
            result["error"] = f"No z_image LoRA for: {', '.join(characters)}"
            return result

    # Build short previz prompt from first_frame + LoRA trigger
    # (Z-Image Turbo only follows ~75 words — full compiler produces 200+)
    prompt_text = build_previz_prompt(shot, characters, lora_registry,
                                      prev_context=prev_context)
    import hashlib
    prompt_hash = hashlib.md5(prompt_text.encode()).hexdigest()[:12]
    result["prompt_hash"] = prompt_hash

    if debug_prompts:
        result["prompt"] = prompt_text
        print(f"    [PROMPT] Shot {shot_id}: {prompt_text[:120]}...")

    # ECU: generate at square aspect ratio, then center-crop to 9:16.
    # Z-Image Turbo composes more naturally in square — less passport-photo centering.
    shot_type = shot.get("shot_type", "").upper()
    if shot_type == "ECU":
        gen_width, gen_height = 896, 896
    else:
        gen_width, gen_height = PREVIZ_WIDTH, PREVIZ_HEIGHT

    # Generate via fal.ai
    t0 = time.time()
    fal_result = generate_previz_frame(
        prompt=prompt_text,
        characters=characters,
        lora_registry=lora_registry,
        width=gen_width,
        height=gen_height,
        seed=seed,
    )
    duration_ms = int((time.time() - t0) * 1000)

    # Download result image
    images = fal_result.get("images", [])
    if not images:
        result["status"] = "failed"
        result["error"] = "No images in fal.ai response"
        if tracker:
            tracker.log_generation(
                episode=episode, shot_id=shot_id, stage="previz",
                model="z_image_turbo_lora", success=False,
                resolution=f"{PREVIZ_WIDTH}x{PREVIZ_HEIGHT}",
                loras=len([c for c in characters
                           if c in lora_registry and lora_registry[c].get(PREVIZ_LORA_KEY)]),
                steps=PREVIZ_STEPS, seed=seed, duration_ms=duration_ms,
            )
        return result

    image_url = images[0].get("url", "")
    download_file(image_url, str(output_path))

    # ECU: center-crop square image to 9:16 vertical
    if shot_type == "ECU" and gen_width == gen_height:
        from PIL import Image
        img = Image.open(str(output_path))
        w, h = img.size
        # Crop center vertical strip: target_w = h * 9/16
        target_w = int(h * 9 / 16)
        left = (w - target_w) // 2
        img = img.crop((left, 0, left + target_w, h))
        img.save(str(output_path))

    result["status"] = "generated"
    result["path"] = str(output_path)
    result["duration_ms"] = duration_ms

    # Log cost
    if tracker:
        tracker.log_generation(
            episode=episode, shot_id=shot_id, stage="previz",
            model="z_image_turbo_lora", success=True,
            resolution=f"{PREVIZ_WIDTH}x{PREVIZ_HEIGHT}",
            loras=len([c for c in characters
                       if c in lora_registry and lora_registry[c].get(PREVIZ_LORA_KEY)]),
            steps=PREVIZ_STEPS, seed=seed, duration_ms=duration_ms,
        )

    return result


# ── HTML Review Page ─────────────────────────────────────────────────────

def build_review_html(
    results: List[dict],
    episode: int,
    seed: int,
    output_dir: Path,
    timestamp: str,
) -> str:
    """Build an HTML contact-sheet review page for previz frames.

    Returns path to the written HTML file.
    """
    generated = [r for r in results if r["status"] == "generated"]
    skipped = [r for r in results if r["status"] == "skipped"]
    blocked = [r for r in results if r["status"] == "blocked_lora"]
    failed = [r for r in results if r["status"] == "failed"]

    html_path = output_dir / "previz_review.html"

    # Build shot cards
    cards_html = ""
    for i, r in enumerate(results):
        shot_id = r["shot_id"]
        char_tag = r.get("char_tag", "ENV")
        gen_type = r.get("generation_type", "unknown")
        status = r["status"]

        # Status badge
        qc_data = r.get("qc", {})
        qc_final = qc_data.get("final", "")
        if status == "qc_rejected":
            badge = '<span class="badge badge-fail">QC REJECT</span>'
        elif status == "generated" and qc_final == "auto_pass":
            badge = '<span class="badge badge-ok">QC PASS</span>'
        elif status == "generated" and qc_final == "edge_case":
            badge = '<span class="badge badge-skip">QC EDGE</span>'
        elif status == "generated":
            badge = '<span class="badge badge-ok">GENERATED</span>'
        elif status == "skipped":
            badge = '<span class="badge badge-skip">SKIPPED</span>'
        elif status == "blocked_lora":
            badge = '<span class="badge badge-block">NO LoRA</span>'
        else:
            badge = '<span class="badge badge-fail">FAILED</span>'

        # Image or placeholder
        if status == "generated" and r.get("path"):
            filename = Path(r["path"]).name
            img_tag = (f'<img src="{filename}" alt="Shot {shot_id}" '
                       f'class="thumb" data-index="{i}" onclick="openLightbox({i})">')
        else:
            img_tag = f'<div class="placeholder">{status.upper()}</div>'

        cards_html += f"""
        <div class="card" data-status="{status}">
            {img_tag}
            <div class="card-info">
                <div class="shot-label">S{shot_id:02d}</div>
                <div class="tags">
                    <span class="tag tag-type">{gen_type.replace('_', ' ')[:16]}</span>
                    <span class="tag tag-char">{char_tag}</span>
                </div>
                {badge}
            </div>
        </div>"""

    # Build lightbox image array (only generated shots)
    lightbox_data = []
    for i, r in enumerate(results):
        if r["status"] == "generated" and r.get("path"):
            lightbox_data.append({
                "index": i,
                "src": Path(r["path"]).name,
                "label": f"S{r['shot_id']:02d} | {r.get('char_tag', 'ENV')}",
            })

    lightbox_json = json.dumps(lightbox_data)

    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Previz Review — Episode {episode}</title>
<style>
:root {{
    --bg: #1a1a2e;
    --surface: #16213e;
    --card-bg: #0f3460;
    --accent: #e94560;
    --text: #eee;
    --text-dim: #999;
    --ok: #4ade80;
    --skip: #facc15;
    --block: #f97316;
    --fail: #ef4444;
}}
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
body {{ background: var(--bg); color: var(--text); font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; padding: 24px; }}
header {{ text-align: center; margin-bottom: 32px; }}
header h1 {{ font-size: 1.6rem; margin-bottom: 8px; }}
header .meta {{ color: var(--text-dim); font-size: 0.85rem; }}
header .meta span {{ margin: 0 8px; }}
.grid {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(180px, 1fr)); gap: 16px; }}
.card {{ background: var(--card-bg); border-radius: 8px; overflow: hidden; transition: transform 0.15s; }}
.card:hover {{ transform: translateY(-2px); }}
.card img.thumb {{ width: 100%; aspect-ratio: 512/896; object-fit: cover; cursor: pointer; display: block; }}
.card .placeholder {{ width: 100%; aspect-ratio: 512/896; display: flex; align-items: center; justify-content: center; background: var(--surface); color: var(--text-dim); font-size: 0.75rem; letter-spacing: 0.05em; }}
.card-info {{ padding: 8px 10px; }}
.shot-label {{ font-weight: 700; font-size: 0.95rem; margin-bottom: 4px; }}
.tags {{ display: flex; gap: 4px; flex-wrap: wrap; margin-bottom: 6px; }}
.tag {{ font-size: 0.65rem; padding: 2px 6px; border-radius: 3px; text-transform: uppercase; letter-spacing: 0.03em; }}
.tag-type {{ background: #1e3a5f; color: #7dd3fc; }}
.tag-char {{ background: #3b1f4e; color: #d8b4fe; }}
.badge {{ font-size: 0.65rem; padding: 2px 8px; border-radius: 10px; font-weight: 600; }}
.badge-ok {{ background: var(--ok); color: #000; }}
.badge-skip {{ background: var(--skip); color: #000; }}
.badge-block {{ background: var(--block); color: #000; }}
.badge-fail {{ background: var(--fail); color: #fff; }}
footer {{ text-align: center; margin-top: 32px; padding-top: 16px; border-top: 1px solid var(--surface); color: var(--text-dim); font-size: 0.8rem; }}
footer .stats span {{ margin: 0 6px; }}

/* Lightbox */
.lightbox {{ display: none; position: fixed; inset: 0; background: rgba(0,0,0,0.92); z-index: 1000; align-items: center; justify-content: center; flex-direction: column; }}
.lightbox.active {{ display: flex; }}
.lightbox img {{ max-height: 85vh; max-width: 90vw; border-radius: 4px; }}
.lightbox .lb-label {{ color: var(--text); margin-top: 12px; font-size: 0.9rem; }}
.lightbox .lb-close {{ position: absolute; top: 16px; right: 24px; font-size: 2rem; color: var(--text-dim); cursor: pointer; background: none; border: none; }}
.lightbox .lb-close:hover {{ color: var(--text); }}
.lb-nav {{ position: absolute; top: 50%; transform: translateY(-50%); font-size: 2.5rem; color: var(--text-dim); cursor: pointer; background: none; border: none; padding: 16px; }}
.lb-nav:hover {{ color: var(--text); }}
.lb-prev {{ left: 12px; }}
.lb-next {{ right: 12px; }}
</style>
</head>
<body>

<header>
    <h1>Episode {episode} — Previz Review</h1>
    <div class="meta">
        <span>{timestamp}</span>
        <span>|</span>
        <span>{len(results)} shots</span>
        <span>|</span>
        <span>seed {seed}</span>
        <span>|</span>
        <span>{PREVIZ_WIDTH}x{PREVIZ_HEIGHT} @ {PREVIZ_STEPS} steps</span>
    </div>
</header>

<div class="grid">
{cards_html}
</div>

<footer>
    <div class="stats">
        <span>Generated: {len(generated)}</span>
        <span>Skipped: {len(skipped)}</span>
        <span>Blocked (no LoRA): {len(blocked)}</span>
        <span>Failed: {len(failed)}</span>
    </div>
</footer>

<!-- Lightbox -->
<div class="lightbox" id="lightbox">
    <button class="lb-close" onclick="closeLightbox()">&times;</button>
    <button class="lb-nav lb-prev" onclick="navLightbox(-1)">&#8249;</button>
    <button class="lb-nav lb-next" onclick="navLightbox(1)">&#8250;</button>
    <img id="lb-img" src="" alt="">
    <div class="lb-label" id="lb-label"></div>
</div>

<script>
const lbData = {lightbox_json};
let lbIdx = 0;

function openLightbox(cardIndex) {{
    const entry = lbData.find(d => d.index === cardIndex);
    if (!entry) return;
    lbIdx = lbData.indexOf(entry);
    showLbImage();
    document.getElementById('lightbox').classList.add('active');
}}

function closeLightbox() {{
    document.getElementById('lightbox').classList.remove('active');
}}

function navLightbox(dir) {{
    lbIdx = (lbIdx + dir + lbData.length) % lbData.length;
    showLbImage();
}}

function showLbImage() {{
    if (!lbData.length) return;
    const entry = lbData[lbIdx];
    document.getElementById('lb-img').src = entry.src;
    document.getElementById('lb-label').textContent = entry.label + ' (' + (lbIdx+1) + '/' + lbData.length + ')';
}}

document.addEventListener('keydown', function(e) {{
    const lb = document.getElementById('lightbox');
    if (!lb.classList.contains('active')) return;
    if (e.key === 'Escape') closeLightbox();
    if (e.key === 'ArrowLeft') navLightbox(-1);
    if (e.key === 'ArrowRight') navLightbox(1);
}});
</script>
</body>
</html>"""

    with open(html_path, "w") as f:
        f.write(html)

    return str(html_path)


# ── Main ─────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Generate previz keyframes (hero frame per shot, low-res, fast)"
    )
    parser.add_argument("project_dir", help="Project directory (e.g., leviathan/)")
    parser.add_argument("--episode", "-e", type=int, required=True, help="Episode number")
    parser.add_argument("--shots", help="Shot filter (e.g., '1-5' or '1,3,8')")
    parser.add_argument("--skip-existing", action="store_true",
                        help="Skip shots that already have previz frames")
    parser.add_argument("--changed-only", action="store_true",
                        help="Only generate shots changed since last storyboard version")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show plan and cost estimate without generating")
    parser.add_argument("--seed", type=int, default=42, help="Random seed (default: 42)")
    parser.add_argument("--debug-prompts", action="store_true",
                        help="Print compiled prompts during generation")
    parser.add_argument("--no-html", action="store_true",
                        help="Skip HTML review page generation")
    parser.add_argument("--qc", action="store_true",
                        help="Run visual gate QC after each frame (Gate 1 artifacts + Gate 2 semantic). "
                             "Auto-rejects trigger regen with new seed (max 2 retries). Requires GOOGLE_API_KEY.")
    args = parser.parse_args()

    # ── Resolve project directory ──
    project_dir = Path(args.project_dir).resolve()
    if not project_dir.exists():
        engine_dir = Path(__file__).resolve().parent.parent.parent
        project_dir = engine_dir / args.project_dir
    if not project_dir.exists():
        print(f"ERROR: Project directory not found: {args.project_dir}", file=sys.stderr)
        sys.exit(1)

    ep_str = f"{args.episode:03d}"

    # ── Pre-flight Validation ──

    # [HARD] Storyboard exists
    storyboard_path = project_dir / "storyboards" / f"storyboard_ep_{ep_str}.json"
    if not storyboard_path.exists():
        print(f"ERROR: Storyboard not found: {storyboard_path}", file=sys.stderr)
        sys.exit(1)

    # [HARD] Storyboard parses
    try:
        with open(storyboard_path) as f:
            storyboard = json.load(f)
    except json.JSONDecodeError as e:
        print(f"ERROR: Invalid JSON in {storyboard_path}: {e}", file=sys.stderr)
        sys.exit(1)

    # [HARD] FAL_KEY set (skip for dry-run)
    if not args.dry_run and not os.environ.get("FAL_KEY"):
        print("ERROR: FAL_KEY environment variable not set", file=sys.stderr)
        sys.exit(1)

    # [HARD] QC requires visual_gate + GOOGLE_API_KEY
    gemini_model = None
    if args.qc:
        if not HAS_VISUAL_GATE:
            print("ERROR: --qc requires visual_gate.py (import failed)", file=sys.stderr)
            sys.exit(1)
        if not os.environ.get("GOOGLE_API_KEY"):
            print("ERROR: --qc requires GOOGLE_API_KEY environment variable", file=sys.stderr)
            sys.exit(1)
        gemini_model = get_gemini_model()

    # [WARNING] Breakdown
    breakdown_path = project_dir / "visual" / "breakdown.json"
    if not breakdown_path.exists():
        print("WARNING: No breakdown.json — prompts will lack wardrobe/location data",
              file=sys.stderr)

    # [WARNING] LoRA registry
    lora_registry_path = project_dir / "visual" / "lora_registry.json"
    if not lora_registry_path.exists():
        print("WARNING: No lora_registry.json — generating without character LoRAs",
              file=sys.stderr)

    # ── Load data ──
    breakdown = _load_breakdown(project_dir)
    project_config = _load_project_config(project_dir)
    override_store = OverrideStore(project_dir)
    asset_code = get_asset_code(project_dir)

    # Load LoRA registry and flatten to inference configs
    raw_registry = load_registry(project_dir)
    lora_registry: Dict[str, dict] = {}
    for char_name in raw_registry:
        lora_registry[char_name] = get_inference_config(raw_registry, char_name)

    # ── Get shots ──
    shots = storyboard.get("shots", [])
    if not shots:
        print("ERROR: No shots in storyboard", file=sys.stderr)
        sys.exit(1)

    # Apply --shots filter
    if args.shots:
        wanted = _parse_shots_arg(args.shots)
        shots = [s for s in shots if s["id"] in wanted]

    # Apply --changed-only filter
    if args.changed_only:
        if get_changed_shots is None:
            print("WARNING: storyboard_version not available, generating all shots",
                  file=sys.stderr)
        else:
            try:
                changed_ids = get_changed_shots(str(storyboard_path))
                if changed_ids:
                    shots = [s for s in shots if s["id"] in set(changed_ids)]
                    print(f"  Changed shots: {sorted(changed_ids)}")
                else:
                    print("WARNING: No storyboard versions found, generating all shots",
                          file=sys.stderr)
            except Exception as e:
                print(f"WARNING: Could not detect changed shots ({e}), generating all",
                      file=sys.stderr)

    # [HARD] At least 1 shot after filtering
    if not shots:
        print("ERROR: No shots match the filter criteria", file=sys.stderr)
        sys.exit(1)

    # ── Output directory ──
    output_dir = project_dir / "storyboards" / "assets" / f"ep_{ep_str}" / "previz"
    output_dir.mkdir(parents=True, exist_ok=True)

    # ── Cost tracker ──
    tracker = CostTracker(project_dir)

    # ── Banner ──
    print(f"\n{'=' * 58}")
    print(f"  PREVIZ GENERATION — Episode {args.episode}")
    print(f"{'=' * 58}")
    print(f"  Storyboard: {storyboard_path.name}")
    print(f"  Shots: {len(shots)}")
    print(f"  Resolution: {PREVIZ_WIDTH}x{PREVIZ_HEIGHT} @ {PREVIZ_STEPS} steps")
    print(f"  Seed: {args.seed}")
    print(f"  Output: {output_dir.relative_to(project_dir)}")
    if args.skip_existing:
        print(f"  Mode: skip-existing")
    if args.changed_only:
        print(f"  Mode: changed-only")
    print(f"{'=' * 58}")

    # ── Dry run ──
    if args.dry_run:
        est_cost_per_shot = 0.004  # ~$0.004 for z_image_turbo 512x896
        total_est = est_cost_per_shot * len(shots)
        print(f"\n  DRY RUN — no API calls will be made\n")
        print(f"  Estimated cost: ${total_est:.3f} ({len(shots)} shots × ${est_cost_per_shot:.4f})")
        print()
        for s in shots:
            sid = s["id"]
            chars = detect_characters(s, storyboard)
            gen_type = determine_generation_type(s)
            char_tag = char_tag_from_list(chars)
            has_lora = any(c in lora_registry and lora_registry[c].get(PREVIZ_LORA_KEY)
                          for c in chars) if chars else False
            lora_status = "OK" if has_lora else ("ENV" if not chars else "NO LoRA")
            print(f"    S{sid:02d}  {char_tag:<12}  {gen_type:<22}  LoRA: {lora_status}")
        print()
        sys.exit(0)

    # ── Generation loop ──
    results: List[dict] = []
    stats = {"generated": 0, "skipped": 0, "blocked_lora": 0, "failed": 0,
             "qc_pass": 0, "qc_edge": 0, "qc_reject": 0, "qc_regen": 0}
    total = len(shots)
    prev_context = None  # Track previous shot context for spatial-aware perspective

    for i, shot in enumerate(shots, 1):
        shot_id = shot["id"]
        characters = detect_characters(shot, storyboard)
        char_tag = char_tag_from_list(characters)
        asset_base = build_asset_name(asset_code, args.episode, shot_id, 1, char_tag)
        existing_path = output_dir / f"{asset_base}_previz.png"

        # Skip existing
        if args.skip_existing and existing_path.exists():
            print(f"  [{i}/{total}] S{shot_id:02d} {char_tag} — skipped (exists)")
            results.append({
                "shot_id": shot_id,
                "status": "skipped",
                "path": str(existing_path),
                "characters": characters,
                "char_tag": char_tag,
                "generation_type": determine_generation_type(shot),
            })
            stats["skipped"] += 1
            prev_context = PreviousShotContext(
                shot_type=shot.get("shot_type", ""),
                camera_angle=shot.get("camera_angle", "eye"),
                camera_side=shot.get("spatial", {}).get("camera_side", ""),
                camera_movement=shot.get("camera_movement", "static"),
                screen_direction=shot.get("spatial", {}).get("screen_direction", ""),
                shot_id=shot.get("id", 0),
            )  # Track even skipped shots for spatial context
            continue

        print(f"  [{i}/{total}] S{shot_id:02d} {char_tag} — generating...", end=" ", flush=True)

        # Generation with optional QC regen loop
        current_seed = args.seed
        attempt = 0
        result = None

        while attempt <= QC_MAX_RETRIES:
            try:
                result = process_shot(
                    shot=shot,
                    storyboard=storyboard,
                    output_dir=output_dir,
                    episode=args.episode,
                    asset_code=asset_code,
                    lora_registry=lora_registry,
                    breakdown=breakdown,
                    project_config=project_config,
                    override_store=override_store,
                    seed=current_seed,
                    debug_prompts=args.debug_prompts,
                    tracker=tracker,
                    prev_context=prev_context,
                )
            except Exception as e:
                result = {
                    "shot_id": shot_id,
                    "status": "failed",
                    "error": str(e),
                    "characters": characters,
                    "char_tag": char_tag,
                    "generation_type": determine_generation_type(shot),
                }
                break

            # If not generated (blocked/failed), no point in QC
            if result["status"] != "generated":
                break

            # If no QC requested, accept the frame
            if not args.qc or not gemini_model:
                break

            # ── Run Visual Gate QC ──
            image_path = result.get("path", "")
            if not image_path or not Path(image_path).exists():
                break

            ref_paths, char_desc = resolve_ref_paths(shot, storyboard, project_dir)
            qc_result = run_gates(
                gemini_model, image_path, shot, ref_paths, char_desc,
                tracker=tracker, episode=args.episode, shot_id=shot_id,
            )

            result["qc"] = qc_result
            final = qc_result.get("final", "error")

            if final == "auto_pass":
                stats["qc_pass"] += 1
                break
            elif final == "edge_case":
                stats["qc_edge"] += 1
                # Accept edge cases but flag them
                break
            elif final == "auto_reject" and attempt < QC_MAX_RETRIES:
                # Regen with different seed
                attempt += 1
                stats["qc_regen"] += 1
                reject_reasons = qc_result.get("reject_reasons", [])
                print(f"QC REJECT ({', '.join(reject_reasons)}), regen #{attempt}...", end=" ", flush=True)
                current_seed = args.seed + (attempt * 1000)
            else:
                # Final rejection or error
                stats["qc_reject"] += 1
                result["status"] = "qc_rejected"
                break

        # Record result
        results.append(result)
        status = result["status"]
        if status == "generated":
            ms = result.get("duration_ms", 0)
            qc_tag = ""
            if args.qc and "qc" in result:
                qc_final = result["qc"].get("final", "")
                if qc_final == "auto_pass":
                    qc_tag = " [QC:PASS]"
                elif qc_final == "edge_case":
                    qc_tag = " [QC:EDGE]"
            print(f"OK ({ms}ms){qc_tag}")
            stats["generated"] += 1
        elif status == "qc_rejected":
            print(f"QC REJECTED after {QC_MAX_RETRIES} retries")
            stats["failed"] += 1
        elif status == "blocked_lora":
            print(f"BLOCKED — {result.get('error', 'no LoRA')}")
            stats["blocked_lora"] += 1
        else:
            print(f"FAILED — {result.get('error', 'unknown')}")
            stats["failed"] += 1

        # Track for spatial-aware perspective on next shot
        prev_context = PreviousShotContext(
            shot_type=shot.get("shot_type", ""),
            camera_angle=shot.get("camera_angle", "eye"),
            camera_side=shot.get("spatial", {}).get("camera_side", ""),
            camera_movement=shot.get("camera_movement", "static"),
            screen_direction=shot.get("spatial", {}).get("screen_direction", ""),
            shot_id=shot.get("id", 0),
        )

        # Brief delay between API calls
        if i < total and stats["generated"] > 0:
            time.sleep(1.0)

    # ── Write manifest ──
    timestamp = datetime.now(timezone.utc).isoformat()
    manifest = {
        "episode": args.episode,
        "generated_at": timestamp,
        "seed": args.seed,
        "resolution": f"{PREVIZ_WIDTH}x{PREVIZ_HEIGHT}",
        "steps": PREVIZ_STEPS,
        "endpoint": PREVIZ_ENDPOINT,
        "stats": stats,
        "shots": results,
    }

    manifest_path = output_dir / "previz_manifest.json"
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)
    print(f"\n  Manifest: {manifest_path.relative_to(project_dir)}")

    # ── Build HTML review ──
    if not args.no_html:
        html_path = build_review_html(
            results=results,
            episode=args.episode,
            seed=args.seed,
            output_dir=output_dir,
            timestamp=datetime.now().strftime("%Y-%m-%d %H:%M"),
        )
        print(f"  Review:   {Path(html_path).relative_to(project_dir)}")

    # ── Summary ──
    print(f"\n{'=' * 58}")
    print(f"  PREVIZ COMPLETE — Episode {args.episode}")
    print(f"{'=' * 58}")
    print(f"  Generated: {stats['generated']}")
    print(f"  Skipped:   {stats['skipped']}")
    print(f"  Blocked:   {stats['blocked_lora']}")
    print(f"  Failed:    {stats['failed']}")
    if args.qc:
        print(f"  --- QC ---")
        print(f"  QC Pass:   {stats['qc_pass']}")
        print(f"  QC Edge:   {stats['qc_edge']}")
        print(f"  QC Reject: {stats['qc_reject']}")
        print(f"  QC Regens: {stats['qc_regen']}")
    print(f"{'=' * 58}\n")

    # Exit code
    sys.exit(1 if stats["failed"] > 0 else 0)


if __name__ == "__main__":
    main()