import os
import tempfile

import spaces
import torch
import gradio as gr
from diffusers import DiffusionPipeline

# Load the pipeline once at startup. The Space is a ZeroGPU space, so the
# model weights stay resident and `@spaces.GPU` allocates a worker per call.
print("Loading Z-Image-Turbo pipeline...")
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
pipe.to("cuda")
print("Pipeline loaded!")


def _save_image(image) -> dict:
    """Mirror of `gradio.workflow._save_tmp`: serialize a PIL.Image as a JSON
    pointer the canvas can render. `Workflow.launch()` already adds the
    tempdir to `allowed_paths`, so the /gradio_api/file=… URL resolves."""
    path = os.path.join(
        tempfile.gettempdir(), f"zimage_{os.urandom(8).hex()}.png"
    )
    image.save(path)
    return {
        "path": path,
        "url": f"/gradio_api/file={path}",
        "orig_name": "zimage.png",
        "mime_type": "image/png",
    }


def _estimate_duration(prompt, height, width, num_inference_steps, seed, randomize_seed) -> int:
    """Rough wall-clock estimate (seconds) for one Z-Image-Turbo call.

    ZeroGPU's default per-call duration is 60s. Requesting less than you need
    *raises* queue priority (shorter tasks get scheduled sooner) and — crucially
    for a busy shared Space — frees the GPU slot for the next visitor far
    sooner than holding it for a full minute, so far fewer users hit the
    Space's "reached its GPU limit" rejection. Scaled by pixel count and steps;
    clamped to a small floor/ceiling so a runaway slider can't starve the queue
    or under-budget a big call.

    Signature mirrors the GPU function exactly because @spaces.GPU passes the
    decorated function's inputs straight through to the duration callable.

    See: https://huggingface.co/docs/hub/en/spaces-zerogpu#duration-management
    """
    pixels = max(int(height), 1) * max(int(width), 1)
    # ~0.4s/step at 1024^2, linear-ish in pixels. Big calls still need headroom.
    per_step = 0.4 * (pixels / (1024 * 1024))
    seconds = int(num_inference_steps) * per_step
    return max(20, min(int(seconds) + 15, 120))


def _friendly_gpu_error(err: Exception) -> str:
    """Turn ZeroGPU's terse allocator rejections into a clear, honest message.

    'Space app has reached its GPU limit' is a *Space-level* capacity rejection
    (the shared ZeroGPU pool is saturated), not a per-user quota wall — it
    reproduces regardless of inputs, account tier, or sign-in state. Don't make
    an upgrade claim whose truth we can't pin down, so the message is neutral:
    shared GPU at capacity, retry shortly.
    """
    msg = (str(err) or "").lower()
    capacity_hints = (
        "gpu limit", "reached its gpu limit", "gpu quota", "out of quota",
        "quota", "no gpu", "could not allocate", "gpu is busy", "too many",
        "concurrent",
    )
    if any(h in msg for h in capacity_hints):
        return (
            "⛔ This demo's shared GPU is at capacity right now — it's not a "
            "problem with your prompt or your account. The GPU pool is fully "
            "booked by other users at the moment. Please wait a minute and "
            "retry; demand clears between bursts."
        )
    if "out of memory" in msg or "oom" in msg or "cuda" in msg:
        return (
            "💥 Image generation ran out of GPU memory. Try a smaller "
            "Height/Width or fewer Inference Steps, then retry."
        )
    return (
        "⚠️ Image generation failed. Please try again in a moment — if it "
        "keeps happening, simplify your prompt or lower the resolution."
    )


@spaces.GPU(duration=_estimate_duration)
def _generate_image_gpu(
    prompt: str,
    height: int,
    width: int,
    num_inference_steps: int,
    seed: int,
    randomize_seed: bool,
):
    """The GPU-decorated worker. Runs only under a ZeroGPU allocation; the
    allocator raises *before* this body if no GPU can be granted, which is why
    the rewording lives in the plain `generate_image` wrapper below, not here.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    if randomize_seed:
        seed = torch.randint(0, 2**32 - 1, (1,)).item()

    generator = torch.Generator("cuda").manual_seed(int(seed))
    image = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=0.0,
        generator=generator,
    ).images[0]

    return _save_image(image), int(seed)


def generate_image(
    prompt: str,
    height: int,
    width: int,
    num_inference_steps: int,
    seed: int,
    randomize_seed: bool,
):
    """Workflow-facing wrapper around the GPU worker.

    Bound to the canvas as a `fn` operator node — the workflow calls this
    Python function directly server-side, so the entire pipeline (frontend +
    ZeroGPU) lives in a single Space. This non-GPU wrapper catches rejections
    from the `@spaces.GPU` allocator (which fire before the worker body runs)
    and rewords them into a clear, honest user-facing message.

    Returns (image_dict, seed_used). The image is serialized to a /gradio_api
    file URL so JSON serialization across the fn bridge succeeds; the
    executor's `fromGradioOutput` turns the dict back into an image port value.
    """
    try:
        return _generate_image_gpu(
            prompt, height, width, num_inference_steps, seed, randomize_seed
        )
    except gr.Error:
        # Already a user-facing validation message (e.g. empty prompt) — pass
        # it through unchanged.
        raise
    except Exception as e:
        # Allocator rejection (GPU limit / quota / OOM / etc.) — reword.
        raise gr.Error(_friendly_gpu_error(e)) from e


# The workflow (workflow.json) wires `generate_image` as a `fn` operator:
#   Prompt, Height, Width, Inference Steps, Seed, Randomize Seed ─▶
#   generate_image (fn operator, kind="fn") ─▶ Output Image, Seed Used
#
# On a Space with `hf_oauth: true`, visiting the canvas runs this function
# under a ZeroGPU worker using each visitor's own HF token.
demo = gr.Workflow(
    graph="workflow.json",
    bind={"generate_image": generate_image},
)

if __name__ == "__main__":
    demo.launch()