"""Test NVIDIA NIM FLUX.1 Kontext-dev for consistent frame variant generation."""
import asyncio
import base64
import os
import time
from pathlib import Path

import httpx
from dotenv import load_dotenv

load_dotenv()

NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
KONTEXT_URL = "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.1-kontext-dev"

# Try to find an existing generated image
WORKDIR = Path("workdir/projects/2/images")


async def generate_test_keyframe() -> bytes:
    """Generate a fresh manhwa-style keyframe via HF FLUX.1-schnell."""
    hf_token = os.getenv("HF_TOKEN")
    prompt = (
        "manhwa style, korean webtoon art, sharp lineart, cel-shaded, "
        "close-up shot, young adult male, jet black messy hair, deep crimson eyes, "
        "black high-collar coat, dramatic lighting, vibrant colors, detailed background"
    )
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.post(
            "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell",
            headers={"Authorization": f"Bearer {hf_token}"},
            json={"inputs": prompt},
        )
        resp.raise_for_status()
        return resp.content


async def kontext_edit(image_b64: str, prompt: str) -> bytes:
    """Send an image to Kontext for editing."""
    async with httpx.AsyncClient(timeout=60) as client:
        resp = await client.post(
            KONTEXT_URL,
            headers={
                "Authorization": f"Bearer {NVIDIA_API_KEY}",
                "Content-Type": "application/json",
                "Accept": "application/json",
            },
            json={
                "image": image_b64,
                "prompt": prompt,
                "steps": 25,
                "cfg_scale": 3.5,
            },
        )
        resp.raise_for_status()
        data = resp.json()
        # Response format: {"artifacts": [{"base64": "..."}]} or {"data": [{"b64_json": "..."}]}
        if "artifacts" in data:
            return base64.b64decode(data["artifacts"][0]["base64"])
        elif "data" in data:
            return base64.b64decode(data["data"][0]["b64_json"])
        else:
            # Try direct b64 field
            for key in ("b64_json", "base64", "image"):
                if key in data:
                    return base64.b64decode(data[key])
            raise ValueError(f"Unknown response format: {list(data.keys())}")


async def main():
    if not NVIDIA_API_KEY:
        print("ERROR: NVIDIA_API_KEY not set in .env")
        return

    output_dir = Path("test_kontext_output")
    output_dir.mkdir(exist_ok=True)

    # Step 1: Get a keyframe image
    print("Step 1: Getting keyframe image...")
    keyframe_path = None

    # Try existing images first
    if WORKDIR.exists():
        images = sorted(WORKDIR.glob("*.png"))
        if images:
            keyframe_path = images[0]
            print(f"  Using existing image: {keyframe_path}")

    if keyframe_path:
        with open(keyframe_path, "rb") as f:
            keyframe_bytes = f.read()
    else:
        print("  Generating fresh keyframe via HF FLUX.1-schnell...")
        keyframe_bytes = await generate_test_keyframe()
        keyframe_path = output_dir / "00_keyframe.png"
        with open(keyframe_path, "wb") as f:
            f.write(keyframe_bytes)
        print(f"  Saved keyframe: {keyframe_path}")

    # Save original
    orig_path = output_dir / "00_original.png"
    with open(orig_path, "wb") as f:
        f.write(keyframe_bytes)

    keyframe_b64 = base64.b64encode(keyframe_bytes).decode()

    # Step 2: Generate 3 variants
    edit_prompts = [
        ("01_smile", "same scene, character now has a slight smile"),
        ("02_hand_raise", "same scene, character raises their right hand"),
        ("03_wind", "same scene, wind blowing through hair, clothes rippling"),
    ]

    for filename, prompt in edit_prompts:
        print(f"\nStep 2: Kontext edit - '{prompt}'...")
        t0 = time.time()
        try:
            variant_bytes = await kontext_edit(keyframe_b64, prompt)
            elapsed = time.time() - t0
            out_path = output_dir / f"{filename}.png"
            with open(out_path, "wb") as f:
                f.write(variant_bytes)
            print(f"  OK in {elapsed:.1f}s -> {out_path}")
        except Exception as e:
            elapsed = time.time() - t0
            print(f"  FAILED in {elapsed:.1f}s: {e}")

    print(f"\nDone! Check {output_dir}/ for results.")
    print("Visual check: Are faces/proportions/colors consistent enough for RIFE interpolation?")


if __name__ == "__main__":
    asyncio.run(main())