#!/usr/bin/env python3
"""
Test: Qwen Image Edit 2511 + Multi-Angle LoRA via fal.ai

Generates 10 multi-angle character images from a single hero reference.
Produces an HTML contact sheet for quality review.

v2 — Uses correct numeric API parameters (horizontal_angle, vertical_angle, zoom)
     instead of prompt-based angle control which was silently ignored.

Usage:
    FAL_KEY=xxx python3 test_qwen_multiangle.py

Endpoint: fal-ai/qwen-image-edit-2511-multiple-angles
Cost: ~$0.035/megapixel → ~$0.035/image at 1024x1024
"""

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

# ── Config ──────────────────────────────────────────────────────────

# Full-res hero image (3MB). Falls back to compressed if upload fails.
HERO_IMAGE_FULLRES = Path(__file__).resolve().parent.parent.parent / \
    "leviathan/visual/refs/characters/heroes/Jinx_Hero.jpeg"
HERO_IMAGE_COMPRESSED = Path("/tmp/jinx_hero_compressed.jpeg")

OUTPUT_DIR = Path(__file__).resolve().parent.parent.parent / \
    "leviathan/visual/lora_candidates/JINX/qwen_test_v2"

ENDPOINT = "fal-ai/qwen-image-edit-2511-multiple-angles"

# 10 test angles — numeric parameters for the Multi-Angle LoRA endpoint
# horizontal_angle: 0=front, 45=3/4 right, 90=profile right, 180=back, 270=profile left, 315=3/4 left
# vertical_angle: -30=low angle, 0=eye level, 30=elevated
# zoom: 0=wide (full body), 5=medium, 10=close-up

TEST_ANGLES = [
    {"name": "01_front_eye_medium",               "h": 0,   "v": 0,   "z": 5,  "desc": "Front, eye level, medium"},
    {"name": "02_three_quarter_right_eye_medium",  "h": 45,  "v": 0,   "z": 5,  "desc": "3/4 right, eye level, medium"},
    {"name": "03_three_quarter_left_eye_medium",   "h": 315, "v": 0,   "z": 5,  "desc": "3/4 left, eye level, medium"},
    {"name": "04_profile_right_eye_medium",        "h": 90,  "v": 0,   "z": 5,  "desc": "Profile right, eye level, medium"},
    {"name": "05_profile_left_eye_medium",         "h": 270, "v": 0,   "z": 5,  "desc": "Profile left, eye level, medium"},
    {"name": "06_back_eye_medium",                 "h": 180, "v": 0,   "z": 5,  "desc": "Back view, eye level, medium"},
    {"name": "07_front_eye_closeup",               "h": 0,   "v": 0,   "z": 10, "desc": "Front, eye level, close-up"},
    {"name": "08_front_low_medium",                "h": 0,   "v": -30, "z": 5,  "desc": "Front, low angle, medium"},
    {"name": "09_front_elevated_wide",             "h": 0,   "v": 30,  "z": 0,  "desc": "Front, elevated, wide (full body)"},
    {"name": "10_three_quarter_right_eye_closeup", "h": 45,  "v": 0,   "z": 10, "desc": "3/4 right, eye level, close-up"},
]


def upload_image(fal_key: str, image_path: Path) -> str:
    """Upload image to fal.ai storage and return URL."""
    import fal_client
    url = fal_client.upload_file(str(image_path))
    print(f"  Uploaded: {image_path.name} → {url[:80]}...")
    return url


def generate_image(fal_key: str, image_url: str, horizontal_angle: int,
                   vertical_angle: int, zoom: int) -> dict:
    """Call Qwen Image Edit Multi-Angle endpoint with numeric parameters."""
    import fal_client

    result = fal_client.subscribe(
        ENDPOINT,
        arguments={
            "image_urls": [image_url],
            "horizontal_angle": horizontal_angle,
            "vertical_angle": vertical_angle,
            "zoom": zoom,
            "lora_scale": 0.9,
            "image_size": "square_hd",       # 1024x1024
            "num_inference_steps": 28,
            "guidance_scale": 4.5,           # endpoint default
            "num_images": 1,
            "enable_safety_checker": False,
        },
        with_logs=False,
    )
    return result


def download_image(url: str, output_path: Path):
    """Download image from URL."""
    req = urllib.request.Request(url)
    with urllib.request.urlopen(req) as resp:
        output_path.write_bytes(resp.read())


def build_contact_sheet(output_dir: Path, results: list, hero_image_name: str):
    """Build HTML contact sheet for visual review."""
    html = """<!DOCTYPE html>
<html><head>
<title>Qwen Multi-Angle Test v2 — JINX</title>
<style>
body { background: #1a1a2e; color: #eee; font-family: system-ui; margin: 20px; }
h1 { color: #e94560; }
.info { color: #aaa; margin-bottom: 20px; }
.grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 16px; }
.card { background: #16213e; border-radius: 8px; overflow: hidden; }
.card img { width: 100%; display: block; cursor: pointer; }
.card img:hover { opacity: 0.9; }
.card .meta { padding: 10px; }
.card .name { font-weight: bold; color: #e94560; }
.card .desc { color: #aaa; font-size: 0.85em; margin-top: 4px; }
.card .params { color: #4ecca3; font-size: 0.8em; margin-top: 4px; font-family: monospace; }
.card .timing { color: #0f3460; font-size: 0.75em; margin-top: 4px; }
.hero-ref { max-width: 250px; border-radius: 8px; margin-bottom: 20px; }
.summary { background: #16213e; padding: 15px; border-radius: 8px; margin-bottom: 20px; }
/* Lightbox */
.lightbox { display: none; position: fixed; top: 0; left: 0; width: 100%; height: 100%;
  background: rgba(0,0,0,0.95); z-index: 1000; cursor: pointer;
  justify-content: center; align-items: center; }
.lightbox.active { display: flex; }
.lightbox img { max-width: 90%; max-height: 90%; object-fit: contain; }
</style>
</head><body>
"""
    html += "<h1>Qwen Image Edit 2511 + Multi-Angle LoRA — JINX Test v2</h1>\n"
    html += f'<p class="info">Generated: {datetime.now().strftime("%Y-%m-%d %H:%M")}'
    html += f' | Endpoint: {ENDPOINT} | Image size: square_hd (1024x1024)</p>\n'
    html += f'<p class="info"><b>v2 fix:</b> Uses numeric params (horizontal_angle, vertical_angle, zoom) instead of prompt strings. '
    html += f'Hero: {hero_image_name} | guidance_scale: 4.5 | lora_scale: 0.9</p>\n'
    html += '<p class="info"><b>Reference image:</b></p>\n'
    html += '<img class="hero-ref" src="../../../refs/characters/heroes/Jinx_Hero.jpeg">\n'

    total_time = sum(r.get("time_s", 0) for r in results)
    total_cost = sum(r.get("cost_est", 0) for r in results)
    html += f'<div class="summary">'
    html += f'<b>Total:</b> {len(results)} images | '
    html += f'Time: {total_time:.1f}s ({total_time/len(results):.1f}s avg) | '
    html += f'Est. cost: ${total_cost:.3f} (${total_cost/len(results):.4f}/img)'
    html += '</div>\n'

    html += '<div class="grid">\n'
    for r in results:
        img_name = r["filename"]
        h, v, z = r.get("h", 0), r.get("v", 0), r.get("z", 5)
        html += f'<div class="card">\n'
        html += f'  <img src="{img_name}" onclick="openLightbox(this.src)">\n'
        html += f'  <div class="meta">\n'
        html += f'    <div class="name">{r["name"]}</div>\n'
        html += f'    <div class="desc">{r["desc"]}</div>\n'
        html += f'    <div class="params">h={h}° v={v}° zoom={z}</div>\n'
        html += f'    <div class="timing">{r.get("time_s", 0):.1f}s</div>\n'
        html += f'  </div>\n'
        html += f'</div>\n'
    html += '</div>\n'

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

    sheet_path = output_dir / "qwen_test_v2_review.html"
    sheet_path.write_text(html)
    return sheet_path


def main():
    fal_key = os.environ.get("FAL_KEY")
    if not fal_key:
        print("ERROR: FAL_KEY not set", file=sys.stderr)
        sys.exit(1)

    # Try full-res hero first, fall back to compressed
    if HERO_IMAGE_FULLRES.exists():
        hero_image = HERO_IMAGE_FULLRES
        hero_label = f"full-res ({hero_image.stat().st_size / 1024 / 1024:.1f}MB)"
    elif HERO_IMAGE_COMPRESSED.exists():
        hero_image = HERO_IMAGE_COMPRESSED
        hero_label = f"compressed ({hero_image.stat().st_size / 1024:.0f}KB)"
    else:
        print(f"ERROR: No hero image found at {HERO_IMAGE_FULLRES} or {HERO_IMAGE_COMPRESSED}",
              file=sys.stderr)
        sys.exit(1)

    # Create output dir
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print("=" * 60)
    print("QWEN IMAGE EDIT 2511 + MULTI-ANGLE LoRA TEST v2")
    print("=" * 60)
    print(f"Hero image:  {hero_image.name} ({hero_label})")
    print(f"Output dir:  {OUTPUT_DIR}")
    print(f"Endpoint:    {ENDPOINT}")
    print(f"Test images: {len(TEST_ANGLES)}")
    print(f"API params:  horizontal_angle, vertical_angle, zoom (NUMERIC)")
    print(f"Est. cost:   ~${len(TEST_ANGLES) * 0.035:.2f}")
    print("=" * 60)

    # Upload hero image once
    print("\nUploading hero image...")
    try:
        image_url = upload_image(fal_key, hero_image)
    except Exception as e:
        if hero_image == HERO_IMAGE_FULLRES and HERO_IMAGE_COMPRESSED.exists():
            print(f"  Full-res upload failed ({e}), falling back to compressed...")
            hero_image = HERO_IMAGE_COMPRESSED
            hero_label = f"compressed ({hero_image.stat().st_size / 1024:.0f}KB)"
            image_url = upload_image(fal_key, hero_image)
        else:
            raise

    # Generate all test images
    results = []
    for i, angle in enumerate(TEST_ANGLES):
        print(f"\n[{i+1}/{len(TEST_ANGLES)}] {angle['desc']}")
        print(f"  Params: h={angle['h']}° v={angle['v']}° zoom={angle['z']}")

        t0 = time.time()
        try:
            result = generate_image(
                fal_key, image_url,
                horizontal_angle=angle["h"],
                vertical_angle=angle["v"],
                zoom=angle["z"],
            )
            elapsed = time.time() - t0

            # Download the image
            if result.get("images") and len(result["images"]) > 0:
                img_url = result["images"][0]["url"]
                filename = f"{angle['name']}.png"
                output_path = OUTPUT_DIR / filename
                download_image(img_url, output_path)
                print(f"  Saved: {filename} ({elapsed:.1f}s)")

                results.append({
                    "name": angle["name"],
                    "desc": angle["desc"],
                    "h": angle["h"],
                    "v": angle["v"],
                    "z": angle["z"],
                    "filename": filename,
                    "time_s": elapsed,
                    "cost_est": 0.035,
                    "width": result["images"][0].get("width", 1024),
                    "height": result["images"][0].get("height", 1024),
                })
            else:
                print(f"  WARNING: No image returned")
                results.append({
                    "name": angle["name"],
                    "desc": angle["desc"],
                    "h": angle["h"],
                    "v": angle["v"],
                    "z": angle["z"],
                    "filename": "FAILED",
                    "time_s": elapsed,
                    "cost_est": 0,
                    "error": "No image in response",
                })

        except Exception as e:
            elapsed = time.time() - t0
            print(f"  ERROR: {e}")
            results.append({
                "name": angle["name"],
                "desc": angle["desc"],
                "h": angle["h"],
                "v": angle["v"],
                "z": angle["z"],
                "filename": "FAILED",
                "time_s": elapsed,
                "cost_est": 0,
                "error": str(e),
            })

    # Write manifest
    manifest = {
        "test": "qwen_multiangle_v2",
        "endpoint": ENDPOINT,
        "hero_image": str(hero_image),
        "hero_label": hero_label,
        "api_params": "numeric (horizontal_angle, vertical_angle, zoom)",
        "guidance_scale": 4.5,
        "lora_scale": 0.9,
        "generated_at": datetime.now().isoformat(),
        "results": results,
    }
    manifest_path = OUTPUT_DIR / "test_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2))

    # Build contact sheet
    sheet_path = build_contact_sheet(OUTPUT_DIR, results, hero_image.name)

    # Summary
    success = [r for r in results if r["filename"] != "FAILED"]
    failed = [r for r in results if r["filename"] == "FAILED"]
    total_time = sum(r["time_s"] for r in results)
    total_cost = sum(r["cost_est"] for r in results)

    print("\n" + "=" * 60)
    print("TEST COMPLETE (v2 — numeric params)")
    print("=" * 60)
    print(f"Success: {len(success)}/{len(results)}")
    if failed:
        print(f"Failed:  {len(failed)}")
        for f in failed:
            print(f"  - {f['name']}: {f.get('error', 'unknown')}")
    print(f"Time:    {total_time:.1f}s total ({total_time/len(results):.1f}s avg)")
    print(f"Cost:    ~${total_cost:.3f}")
    print(f"\nReview:  {sheet_path}")
    print(f"Output:  {OUTPUT_DIR}")


if __name__ == "__main__":
    main()