#!/usr/bin/env python3
"""
Pipeline Variant Test — Skip Qwen Edit, test NBP for background+expression directly.

Tests:
1. Hero → Qwen MA (3/4) → NBP (background + expression) → SeedVR2
2. Hero → Qwen MA (close-up) → NBP (different background + expression) → SeedVR2
3. Hero → Qwen MA (low angle) → SeedVR2 (no NBP, just quality)

Usage:
    python3 pipeline_variant_test.py
"""

import os
import sys
import time
import urllib.request
from pathlib import Path
from datetime import datetime

try:
    import fal_client
except ImportError:
    print("ERROR: fal_client not installed")
    sys.exit(1)

# --- Paths ---
ENGINE_DIR = Path(__file__).resolve().parent
PROJECT_DIR = ENGINE_DIR.parent.parent / "leviathan"
SHOOTOUT_DIR = PROJECT_DIR / "visual" / "lora_candidates" / "JINX" / "shootout"
OUTPUT_DIR = SHOOTOUT_DIR / f"pipeline_variant_test_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

# Import angle maps from engine_shootout
sys.path.insert(0, str(ENGINE_DIR))
from engine_shootout import (
    QWEN_ANGLE_MAP, ANGLE_DESCRIPTIONS,
    _resolve_hero,
)

# --- Test cases ---
TESTS = [
    {
        "name": "three_quarter_right_exhausted",
        "label": "3/4 Right — Exhausted",
        "angle": "three_quarter_right",
        "expression": "exhausted — heavy-lidded eyes, slight frown, drained hollow gaze",
        "environment": "dimly lit industrial corridor with exposed pipes and warm overhead sodium light",
        "use_nbp": True,
    },
    {
        "name": "closeup_front_focused",
        "label": "Close-up Front — Focused",
        "angle": "closeup_front",
        "expression": "focused — narrowed eyes, set jaw, intent forward stare",
        "environment": "rain-slicked city rooftop at dusk with neon reflections and cool blue ambient light",
        "use_nbp": True,
    },
    {
        "name": "low_angle_neutral",
        "label": "Low Angle — Neutral (no NBP)",
        "angle": "low_angle",
        "expression": "neutral — relaxed features, steady gaze, lips closed naturally",
        "environment": None,  # Not needed — no NBP pass
        "use_nbp": False,
    },
]


def upload_to_fal(image_path: Path, retries: int = 5) -> str:
    """Upload a local image to fal.ai with retry."""
    data = image_path.read_bytes()
    mime = "image/jpeg" if image_path.suffix.lower() in (".jpg", ".jpeg") else "image/png"
    for attempt in range(retries):
        try:
            url = fal_client.upload(data, mime, file_name="input.png")
            return url
        except Exception as e:
            if attempt < retries - 1:
                wait = 10 * (attempt + 1)
                print(f"  Upload failed (attempt {attempt+1}/{retries}), retrying in {wait}s...")
                time.sleep(wait)
            else:
                raise


def run_qwen_ma(hero_path: Path, output_path: Path, angle: str) -> dict:
    """Pass 1: Qwen Multi-Angle."""
    print(f"  [Pass 1] Qwen MA → {angle}...")
    image_url = upload_to_fal(hero_path)
    angle_params = QWEN_ANGLE_MAP.get(angle, {"h": 315, "v": 0, "z": 5})

    t0 = time.time()
    result = fal_client.subscribe(
        "fal-ai/qwen-image-edit-2511-multiple-angles",
        arguments={
            "image_urls": [image_url],
            "horizontal_angle": angle_params["h"],
            "vertical_angle": angle_params["v"],
            "zoom": angle_params["z"],
            "lora_scale": 0.9,
            "image_size": "square_hd",
            "num_inference_steps": 28,
            "guidance_scale": 4.5,
            "num_images": 1,
            "enable_safety_checker": False,
        },
        with_logs=False,
    )
    elapsed = time.time() - t0

    if result and "images" in result and result["images"]:
        img_url = result["images"][0]["url"]
        output_path.parent.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(img_url, str(output_path))
        print(f"  [Pass 1] Done ({elapsed:.1f}s)")
        return {"success": True, "elapsed": elapsed, "output": output_path}
    print(f"  [Pass 1] FAILED ({elapsed:.1f}s)")
    return {"success": False, "elapsed": elapsed}


def run_nbp_background_expression(input_path: Path, output_path: Path,
                                   angle: str, expression: str,
                                   environment: str, lighting: str = None) -> dict:
    """NBP for background swap + expression (NO dual-reference, single input only)."""
    from google import genai
    from google.genai import types

    print(f"  [Pass 2] NBP → background + expression...")
    client = genai.Client(api_key=os.environ.get("GOOGLE_API_KEY"))

    img_bytes = input_path.read_bytes()
    mime = "image/jpeg" if input_path.suffix.lower() in (".jpg", ".jpeg") else "image/png"
    angle_desc = ANGLE_DESCRIPTIONS.get(angle, "3/4 view, eye level")
    lighting_desc = lighting or "cinematic lighting with modeling"

    prompt_text = (
        "This is the SAME PERSON — face identity locked. DO NOT generate a new face.\n\n"
        f"Change the background/environment to: {environment}.\n"
        f"The subject's expression: {expression}.\n"
        f"Camera angle: {angle_desc}.\n"
        f"Lighting: {lighting_desc}.\n\n"
        "Maintain identical skull structure, brow ridge, nose bridge, nose width, "
        "cheekbone position, chin shape, ear shape, eye spacing, eye size, iris color, "
        "skin tone, skin texture, hair color, and hair texture.\n"
        "Expression muscles may move freely — brows, eyelids, nostrils, lips, jaw — "
        "these are temporary muscular movements, not identity changes.\n\n"
        "Keep the EXACT same camera angle, head position, and body orientation. "
        "DO NOT rotate the head. DO NOT change the framing or crop.\n\n"
        "Shot at 85mm f/1.8. Sharp focus on eyes. Tack-sharp iris detail with visible iris fibers. "
        "Subtle catchlights in the pupils.\n\n"
        "DO NOT smooth, beautify, or stylize facial features. No global smoothing. No airbrushing.\n"
        "Strictly preserve pore texture, freckles, fine lines, and skin imperfections.\n"
        "Single photorealistic photograph. One person only. No text. No split panels."
    )

    parts = [
        types.Part(inline_data=types.Blob(mime_type=mime, data=img_bytes)),
        types.Part(text=prompt_text),
    ]
    contents = [types.Content(parts=parts)]

    t0 = time.time()
    response = client.models.generate_content(
        model="gemini-3-pro-image-preview",
        contents=contents,
        config=types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"],
            image_config=types.ImageConfig(aspect_ratio="1:1"),
            temperature=1.0,
        ),
    )
    elapsed = time.time() - t0

    # Extract image
    for part in (response.candidates[0].content.parts if response.candidates else []):
        if hasattr(part, "inline_data") and part.inline_data and part.inline_data.mime_type.startswith("image/"):
            import base64
            img_data = part.inline_data.data
            if isinstance(img_data, str):
                img_data = base64.b64decode(img_data)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.write_bytes(img_data)
            print(f"  [Pass 2] Done ({elapsed:.1f}s)")
            return {"success": True, "elapsed": elapsed, "output": output_path}

    print(f"  [Pass 2] FAILED — no image in response ({elapsed:.1f}s)")
    return {"success": False, "elapsed": elapsed}


def run_seedvr2(input_path: Path, output_path: Path) -> dict:
    """SeedVR2 quality upscale."""
    print(f"  [Quality] SeedVR2 upscale...")
    image_url = upload_to_fal(input_path)

    t0 = time.time()
    result = fal_client.run(
        "fal-ai/seedvr/upscale/image",
        arguments={"image_url": image_url},
    )
    elapsed = time.time() - t0

    if result and "image" in result:
        img_url = result["image"]["url"]
        output_path.parent.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(img_url, str(output_path))
        print(f"  [Quality] Done ({elapsed:.1f}s)")
        return {"success": True, "elapsed": elapsed, "output": output_path}
    print(f"  [Quality] FAILED ({elapsed:.1f}s)")
    return {"success": False, "elapsed": elapsed}


def generate_html(test_results: list, output_dir: Path) -> Path:
    """Generate comparison HTML."""
    html = """<!DOCTYPE html>
<html>
<head>
<title>Pipeline Variant Test — Skip Qwen Edit</title>
<style>
    body { background: #1a1a1a; color: #e0e0e0; font-family: -apple-system, sans-serif; margin: 20px; }
    h1 { color: #fff; border-bottom: 2px solid #444; padding-bottom: 10px; }
    h2 { color: #aaa; margin-top: 40px; }
    .note { color: #999; font-size: 14px; margin: 10px 0; padding: 10px; background: #222; border-radius: 4px; }
    .test-case { margin-bottom: 60px; }
    .section-label { font-size: 16px; color: #0af; margin: 20px 0 8px 0; font-weight: bold; }
    .row { display: flex; gap: 12px; margin-bottom: 20px; flex-wrap: wrap; }
    .col { flex: 1; min-width: 250px; max-width: 520px; }
    .col img { width: 100%; border: 2px solid #333; border-radius: 4px; cursor: pointer; }
    .col img:hover { border-color: #0af; }
    .label { text-align: center; font-size: 13px; color: #888; margin-top: 4px; }
    .label strong { color: #ccc; }
    .time { color: #666; font-size: 11px; }
    .pipeline { color: #0f0; font-size: 12px; font-family: monospace; }
    .lightbox { display: none; position: fixed; top: 0; left: 0; width: 100vw; height: 100vh;
                background: rgba(0,0,0,0.95); z-index: 1000; justify-content: center; align-items: center; }
    .lightbox.active { display: flex; }
    .lightbox img { max-width: 95vw; max-height: 95vh; object-fit: contain; }
    .lightbox-close { position: fixed; top: 20px; right: 30px; color: #fff; font-size: 30px; cursor: pointer; z-index: 1001; }
</style>
</head>
<body>
<h1>Pipeline Variant Test — Skip Qwen Edit</h1>
<p class="note">
<strong>Hypothesis:</strong> Qwen Edit (Pass 2) may be the weak link. Testing NBP directly on Qwen MA output for background + expression in a single pass.<br>
<strong>Pipelines tested:</strong><br>
<span class="pipeline">Test 1 & 2: Hero &rarr; Qwen MA &rarr; NBP (bg + expression) &rarr; SeedVR2</span><br>
<span class="pipeline">Test 3: Hero &rarr; Qwen MA &rarr; SeedVR2 (no NBP)</span><br>
Click any image to enlarge.
</p>
"""

    for tc in test_results:
        html += f'<div class="test-case">\n'
        html += f'<h2>{tc["label"]}</h2>\n'
        html += f'<p class="pipeline">{tc["pipeline_desc"]}</p>\n'

        html += '<div class="row">\n'
        for step in tc["steps"]:
            html += f'<div class="col"><img src="{step["file"]}" onclick="openLightbox(this)">'
            html += f'<div class="label"><strong>{step["label"]}</strong><br>{step["desc"]}'
            if step.get("time"):
                html += f'<br><span class="time">{step["time"]}</span>'
            html += '</div></div>\n'
        html += '</div>\n'  # row
        html += '</div>\n'  # test-case

    html += """
<div id="lightbox" class="lightbox" onclick="closeLightbox()">
    <span class="lightbox-close">&times;</span>
    <img id="lightbox-img" src="">
</div>
<script>
function openLightbox(el) {
    document.getElementById('lightbox-img').src = el.src;
    document.getElementById('lightbox').classList.add('active');
}
function closeLightbox() {
    document.getElementById('lightbox').classList.remove('active');
}
document.addEventListener('keydown', e => { if (e.key === 'Escape') closeLightbox(); });
</script>
</body></html>"""

    html_path = output_dir / "comparison.html"
    html_path.write_text(html)
    return html_path


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    # Find hero
    hero_path = _resolve_hero(PROJECT_DIR, "JINX")
    if not hero_path or not hero_path.exists():
        print("ERROR: Hero image not found")
        sys.exit(1)
    print(f"Hero: {hero_path.name}")

    test_results = []

    for tc in TESTS:
        print(f"\n{'='*60}")
        print(f"TEST: {tc['label']}")
        print(f"{'='*60}")

        steps = []

        # Step 1: Qwen MA (skip if already exists)
        qwen_out = OUTPUT_DIR / f"{tc['name']}_pass1_qwen_ma.png"
        if qwen_out.exists():
            print(f"  [Pass 1] Already exists, skipping")
            qwen_result = {"success": True, "elapsed": 0}
        else:
            qwen_result = run_qwen_ma(hero_path, qwen_out, tc["angle"])
        if not qwen_result["success"]:
            print(f"  SKIPPING — Qwen MA failed")
            continue
        steps.append({
            "file": qwen_out.name,
            "label": "Pass 1: Qwen MA",
            "desc": f"Angle: {tc['angle']}",
            "time": f"{qwen_result['elapsed']:.1f}s",
        })

        if tc["use_nbp"]:
            # Step 2: NBP (background + expression, single input, no dual-ref)
            nbp_out = OUTPUT_DIR / f"{tc['name']}_pass2_nbp.png"
            if nbp_out.exists():
                print(f"  [Pass 2] Already exists, skipping")
                nbp_result = {"success": True, "elapsed": 0}
            else:
                nbp_result = run_nbp_background_expression(
                    qwen_out, nbp_out,
                    angle=tc["angle"],
                    expression=tc["expression"],
                    environment=tc["environment"],
                )
            if not nbp_result["success"]:
                print(f"  NBP failed — continuing without it")
                # Still run SeedVR2 on Qwen MA output
                seed_input = qwen_out
            else:
                steps.append({
                    "file": nbp_out.name,
                    "label": "Pass 2: NBP",
                    "desc": f"Background + expression (no Qwen Edit)",
                    "time": f"{nbp_result['elapsed']:.1f}s",
                })
                seed_input = nbp_out

            # Step 3: SeedVR2
            seed_out = OUTPUT_DIR / f"{tc['name']}_pass3_seedvr2.png"
            if seed_out.exists():
                print(f"  [Quality] Already exists, skipping")
                seed_result = {"success": True, "elapsed": 0}
            else:
                seed_result = run_seedvr2(seed_input, seed_out)
            if seed_result["success"]:
                steps.append({
                    "file": seed_out.name,
                    "label": "Pass 3: SeedVR2",
                    "desc": "Non-generative quality upscale",
                    "time": f"{seed_result['elapsed']:.1f}s",
                })

            pipeline_desc = "Hero → Qwen MA → NBP (bg + expr) → SeedVR2"
        else:
            # No NBP — just SeedVR2
            seed_out = OUTPUT_DIR / f"{tc['name']}_pass2_seedvr2.png"
            if seed_out.exists():
                print(f"  [Quality] Already exists, skipping")
                seed_result = {"success": True, "elapsed": 0}
            else:
                seed_result = run_seedvr2(qwen_out, seed_out)
            if seed_result["success"]:
                steps.append({
                    "file": seed_out.name,
                    "label": "Pass 2: SeedVR2",
                    "desc": "Non-generative quality upscale (no NBP)",
                    "time": f"{seed_result['elapsed']:.1f}s",
                })
            pipeline_desc = "Hero → Qwen MA → SeedVR2 (no NBP)"

        test_results.append({
            "name": tc["name"],
            "label": tc["label"],
            "pipeline_desc": pipeline_desc,
            "steps": steps,
        })

    if test_results:
        html_path = generate_html(test_results, OUTPUT_DIR)
        print(f"\n{'='*60}")
        print(f"ALL TESTS COMPLETE")
        print(f"{'='*60}")
        print(f"  Output: {OUTPUT_DIR}")
        print(f"  Comparison: {html_path}")

        # Open in browser
        os.system(f'open "{html_path}"')


if __name__ == "__main__":
    main()