import os
import random

import spaces
import torch
from diffusers import QwenImagePipeline
from fastapi.responses import HTMLResponse
from gradio import Server
from gradio.data_classes import FileData

torch.set_float32_matmul_precision("high")

MODEL_ID = "nvidia/Qwen-Image-Flash"
MAX_SEED = 2**31 - 1

# Load the pipeline at module top, *not* inside @spaces.GPU. Reasons:
# 1. ZeroGPU docs explicitly say lazy-loading inside @spaces.GPU is
#    discouraged — module-level placement uses PyTorch's CUDA emulation
#    mode outside the decorator, so `pipe.to("cuda")` works on the
#    CPU-only Space builder, and the real-CUDA transfers are optimized
#    for placements done at startup rather than inside a handler.
# 2. Crucially, `from_pretrained` runs during *module import* — which the
#    Space builder executes as part of the build stage, BEFORE the runtime
#    container reports "Running." That means the ~20 GB checkpoint lands
#    in the cache during build, so the runtime container starts up with
#    the weights already on disk and the first request can begin GPU work
#    immediately rather than downloading first.
pipe = QwenImagePipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
pipe.to("cuda")


def _get_duration(
    prompt: str,
    negative_prompt: str | None,
    width: int,
    height: int,
    num_inference_steps: int,
    seed: int,
    randomize_seed: bool,
) -> int:
    """Return the GPU time budget the ZeroGPU scheduler should grant.

    4 denoising steps over 1024^2 is ~30-60 s on Blackwell. We size the budget
    loosely around `steps * 15s + 30s` (clamped 60-180 s) so larger step counts
    and resolutions don't get cut off, but stay under the 3-min ceiling where
    queue priority starts to suffer.
    """
    steps = max(1, int(num_inference_steps))
    return min(180, max(60, steps * 15 + 30))


app = Server()


@app.api
@spaces.GPU(size="xlarge", duration=_get_duration)
def generate_image(
    prompt: str,
    negative_prompt: str | None,
    width: int,
    height: int,
    num_inference_steps: int,
    seed: int,
    randomize_seed: bool,
) -> tuple[FileData, int, str]:
    """Generate an image from a text prompt with Qwen-Image-Flash.

    The DMD2 student internalized CFG=4.0 during distillation, so inference runs
    at ``true_cfg_scale=1.0`` with ``guidance_scale=None`` / ``negative_prompt=None``
    to avoid applying guidance a second time. The shift-3 FlowMatch Euler
    scheduler produces the required 4-step trajectory. Width and height must be
    divisible by 16; ``1024 x 1024`` is the tested setting.
    """
    if not prompt or not prompt.strip():
        raise ValueError("Prompt cannot be empty.")

    # Snap to a multiple of 16 (the VAE constraint) and clamp to a sane range.
    width = max(256, min(2048, (int(width) // 16) * 16))
    height = max(256, min(2048, (int(height) // 16) * 16))

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed) % (MAX_SEED + 1)

    negative = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None

    generator = torch.Generator(device="cuda").manual_seed(seed)

    result = pipe(
        prompt=prompt,
        negative_prompt=negative,
        width=width,
        height=height,
        num_inference_steps=int(num_inference_steps),
        true_cfg_scale=1.0,
        guidance_scale=None,
        generator=generator,
    )
    image = result.images[0]

    # Persist to disk so the queue can ship it back to the JS client as a FileData URL.
    out_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"qwen-image-flash-{seed}-{width}x{height}.png")
    image.save(out_path)

    info = (
        f"seed={seed}  |  {width}x{height}  |  {num_inference_steps} steps  |  "
        f"true_cfg_scale=1.0"
    )
    return FileData(path=out_path), seed, info


@app.get("/")
async def homepage() -> HTMLResponse:
    """Serve the custom vanilla-HTML frontend."""
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read(), media_type="text/html")


if __name__ == "__main__":
    app.launch(show_error=True)