import spaces  # must precede torch / diffusers

import os
import random
import tempfile
import time

import gradio as gr
import torch
from PIL import Image

from diffusers import QwenImage21Pipeline


MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen-Image-2.1")
MAX_SEED = 2**31 - 1
MAX_REFERENCE_IMAGES = 10
MAX_DURATION = 195

# Qwen Image's latent-space geometry uses 16x spatial compression.
SIZE_MULTIPLE = 16
MAX_EDIT_LONG_SIDE = 2048


# Text-to-image output presets.
SIZE_PRESETS = {
    "1:1 — 1024x1024": (1024, 1024),
    "16:9 — 1376x768": (1376, 768),
    "9:16 — 768x1376": (768, 1376),
    "4:3 — 1200x896": (1200, 896),
    "3:4 — 896x1200": (896, 1200),
    "3:2 — 1264x848": (1264, 848),
    "2:3 — 848x1264": (848, 1264),
    "1:1 — 2048x2048 (2K)": (2048, 2048),
    "16:9 — 2752x1536 (2K)": (2752, 1536),
    "9:16 — 1536x2752 (2K)": (1536, 2752),
    "4:3 — 2400x1792 (2K)": (2400, 1792),
    "3:4 — 1792x2400 (2K)": (1792, 2400),
    "3:2 — 2528x1696 (2K)": (2528, 1696),
    "2:3 — 1696x2528 (2K)": (1696, 2528),
}
DEFAULT_PRESET = "1:1 — 2048x2048 (2K)"


# Image-to-image output settings.
EDIT_SIZE_PRESETS = {
    "2048x2048": "square_2k",
    "2048": "long_side_2k",
}
DEFAULT_EDIT_SIZE = "2048x2048"


# The model card's recommended wrapper for native transparent output.
RGBA_PREFIX = "This is an RGBA image with transparency. "
RGBA_SUFFIX = " The image has alpha channel and the background is transparent."


pipe = QwenImage21Pipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to("cuda")


# Avoid a large full-frame VAE decode. At 1024px, ordinary 1K images still decode
# in one tile, while larger outputs use tiled decoding and require less peak VRAM.
pipe.vae.enable_tiling(
    tile_sample_min_height=1024,
    tile_sample_min_width=1024,
    tile_sample_stride_height=768,
    tile_sample_stride_width=768,
)


def _resolve_seed(seed: int, randomize: bool) -> int:
    if randomize or seed is None or int(seed) < 0:
        return random.randint(0, MAX_SEED)
    return int(seed) % (MAX_SEED + 1)


def _save_png(image: Image.Image) -> str:
    """Save an image as a unique PNG file and return its path."""
    fd, path = tempfile.mkstemp(suffix=".png")
    os.close(fd)
    image.save(path, format="PNG")
    return path


def _round_to_multiple(value: float, multiple: int = SIZE_MULTIPLE) -> int:
    """Round a dimension to the nearest valid Qwen latent-grid multiple."""
    return max(multiple, int(round(value / multiple)) * multiple)


def _edit_output_size(
    first_image: Image.Image,
    edit_size: str = DEFAULT_EDIT_SIZE,
) -> tuple[int, int]:
    """Determine the target canvas for an image-edit request.

    Square mode always produces 2048x2048.

    Long-side mode preserves the first reference image's aspect ratio, scales its
    longest side to 2048 px, and rounds both dimensions to multiples of 16.
    """
    mode = EDIT_SIZE_PRESETS.get(
        edit_size,
        EDIT_SIZE_PRESETS[DEFAULT_EDIT_SIZE],
    )

    if mode == "square_2k":
        return 2048, 2048

    source_width, source_height = first_image.size
    if source_width <= 0 or source_height <= 0:
        raise gr.Error("The first reference image has invalid dimensions.")

    scale = MAX_EDIT_LONG_SIDE / max(source_width, source_height)

    width = _round_to_multiple(source_width * scale)
    height = _round_to_multiple(source_height * scale)

    # Rounding upward can make an edge 16 pixels too large. Scale down once more
    # so that "long side 2048" always remains a real upper bound.
    if max(width, height) > MAX_EDIT_LONG_SIDE:
        second_scale = MAX_EDIT_LONG_SIDE / max(width, height)
        width = _round_to_multiple(width * second_scale)
        height = _round_to_multiple(height * second_scale)

    # This is a final guard against edge cases from rounding.
    width = min(width, MAX_EDIT_LONG_SIDE)
    height = min(height, MAX_EDIT_LONG_SIDE)

    return width, height


def _tokens(width: int, height: int) -> int:
    """Number of spatial latent tokens after 16x compression."""
    return (width // 16) * (height // 16)


def _raw_budget(
    width: int,
    height: int,
    steps: int,
    references: int = 0,
) -> int:
    """Estimate the un-clamped ZeroGPU duration requirement.

    The output latents form attention queries. Reference-image latents are
    key/value context and are cached after the first step. This makes the
    approximate attention cost O x (O + R), not (O + R)^2.
    """
    out_tokens = _tokens(width, height)
    ref_tokens = out_tokens * references

    per_step = (
        9.11e-5 * out_tokens
        + 5.78e-9 * out_tokens * (out_tokens + ref_tokens)
    )
    seconds = 12 + 4 * references + steps * per_step

    return int(max(40, seconds * 1.3))


def _budget(
    width: int,
    height: int,
    steps: int,
    references: int = 0,
) -> int:
    """Return a reservation duration no larger than ZeroGPU permits."""
    return min(
        MAX_DURATION,
        _raw_budget(width, height, steps, references),
    )


def _require_budget(
    width: int,
    height: int,
    steps: int,
    references: int = 0,
) -> None:
    """Reject work that cannot fit inside a single allowed ZeroGPU reservation."""
    required = _raw_budget(width, height, steps, references)

    if required > MAX_DURATION:
        raise gr.Error(
            "This request needs more GPU time than a single ZeroGPU reservation "
            "allows. Use fewer steps, fewer reference images, or a smaller size."
        )


def _gallery_item_path(item) -> str:
    """Extract a filepath from Gradio Gallery's supported return formats."""
    if isinstance(item, (tuple, list)):
        return item[0]
    return item


def _t2i_duration(
    prompt=None,
    transparent=False,
    size_preset=DEFAULT_PRESET,
    steps=30,
    *args,
    **kwargs,
):
    width, height = SIZE_PRESETS.get(
        size_preset,
        SIZE_PRESETS[DEFAULT_PRESET],
    )
    return _budget(width, height, int(steps or 30))


def _edit_duration(
    instruction=None,
    reference_images=None,
    edit_size=DEFAULT_EDIT_SIZE,
    steps=30,
    *args,
    **kwargs,
):
    """Estimate image-edit duration from the actual selected output geometry."""
    references = reference_images or []
    count = len(references)

    # Before an image is uploaded, reserve enough for the selected 2K operation.
    if not references:
        return _budget(
            2048,
            2048,
            int(steps or 30),
            references=0,
        )

    try:
        first_path = _gallery_item_path(references[0])

        with Image.open(first_path) as first_image:
            width, height = _edit_output_size(first_image, edit_size)

        return _budget(
            width,
            height,
            int(steps or 30),
            references=count,
        )
    except Exception:
        # The browser may be mid-upload when the duration callback is evaluated.
        # Use a safe upper-bound fallback in that situation.
        return _budget(
            2048,
            2048,
            int(steps or 30),
            references=count,
        )


@spaces.GPU(duration=_t2i_duration)
def generate(
    prompt: str,
    transparent: bool = False,
    size_preset: str = DEFAULT_PRESET,
    steps: int = 30,
    true_cfg_scale: float = 1.0,
    negative_prompt: str = "",
    seed: int = 0,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image from a text prompt."""
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt first.")

    text = prompt.strip()
    if transparent:
        text = RGBA_PREFIX + text + RGBA_SUFFIX

    width, height = SIZE_PRESETS.get(
        size_preset,
        SIZE_PRESETS[DEFAULT_PRESET],
    )

    _require_budget(width, height, int(steps))
    used_seed = _resolve_seed(seed, randomize_seed)

    kwargs = dict(
        prompt=text,
        width=width,
        height=height,
        num_inference_steps=int(steps),
        true_cfg_scale=float(true_cfg_scale),
        generator=torch.Generator("cuda").manual_seed(used_seed),
    )

    if (
        float(true_cfg_scale) > 1.0
        and negative_prompt
        and negative_prompt.strip()
    ):
        kwargs["negative_prompt"] = negative_prompt.strip()

    started = time.perf_counter()
    image = pipe(**kwargs).images[0]
    elapsed = time.perf_counter() - started

    # An opaque text-to-image result is exported as RGB PNG. Native transparent
    # results retain their RGBA alpha channel.
    if not transparent and image.mode == "RGBA":
        image = image.convert("RGB")

    return (
        image,
        _save_png(image),
        used_seed,
        f"{elapsed:.1f}s on {width}x{height}, {int(steps)} steps",
    )


@spaces.GPU(duration=_edit_duration)
def edit(
    instruction: str,
    reference_images: list,
    edit_size: str = DEFAULT_EDIT_SIZE,
    steps: int = 30,
    true_cfg_scale: float = 1.0,
    negative_prompt: str = "",
    seed: int = 0,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit or recombine up to ten images using a natural-language instruction."""
    if not instruction or not instruction.strip():
        raise gr.Error("Describe the edit you want.")

    if not reference_images:
        raise gr.Error("Add at least one reference image.")

    if len(reference_images) > MAX_REFERENCE_IMAGES:
        raise gr.Error(
            f"Up to {MAX_REFERENCE_IMAGES} reference images are supported."
        )

    images = []

    for item in reference_images:
        path = _gallery_item_path(item)

        # Copy while the input file is open. This closes file handles before GPU
        # inference and avoids lazy reads from temporary Gradio upload paths.
        with Image.open(path) as source:
            images.append(source.copy())

    # The first reference defines the aspect ratio for long-side mode.
    width, height = _edit_output_size(images[0], edit_size)

    _require_budget(
        width,
        height,
        int(steps),
        references=len(images),
    )

    used_seed = _resolve_seed(seed, randomize_seed)

    kwargs = dict(
        prompt=instruction.strip(),
        image=images,
        width=width,
        height=height,
        num_inference_steps=int(steps),
        true_cfg_scale=float(true_cfg_scale),
        generator=torch.Generator("cuda").manual_seed(used_seed),
    )

    if (
        float(true_cfg_scale) > 1.0
        and negative_prompt
        and negative_prompt.strip()
    ):
        kwargs["negative_prompt"] = negative_prompt.strip()

    started = time.perf_counter()
    image = pipe(**kwargs).images[0]
    elapsed = time.perf_counter() - started

    # _save_png preserves RGBA alpha when the pipeline supplies it.
    return (
        image,
        _save_png(image),
        used_seed,
        f"{elapsed:.1f}s from {len(images)} reference image(s), "
        f"{width}x{height}, {int(steps)} steps",
    )


INTRO = """ """


with gr.Blocks(title=" ") as demo:
    gr.Markdown(INTRO)

    with gr.Tabs():
        with gr.Tab(" "):
            with gr.Row():
                with gr.Column(scale=1):
                    t2i_prompt = gr.Textbox(
                        label=" ",
                        lines=10,
                    )

                    t2i_transparent = gr.Checkbox(
                        label=" ",
                        value=False,
                    )

                    t2i_size = gr.Dropdown(
                        label=" ",
                        choices=list(SIZE_PRESETS),
                        value=DEFAULT_PRESET,
                    )

                    t2i_run = gr.Button(
                        " ",
                        variant="primary",
                    )

                    with gr.Accordion(" ", open=False):
                        t2i_steps = gr.Slider(
                            10,
                            50,
                            value=25,
                            step=1,
                            label=" ",
                        )

                        t2i_cfg = gr.Slider(
                            1.0,
                            10.0,
                            value=1.1,
                            step=0.1,
                            label=" ",
                        )

                        t2i_negative = gr.Textbox(
                            label=" ",
                            lines=2,
                        )

                        t2i_seed = gr.Slider(
                            0,
                            MAX_SEED,
                            value=0,
                            step=1,
                            label=" ",
                        )

                        t2i_random = gr.Checkbox(
                            label=" ",
                            value=True,
                        )

                with gr.Column(scale=1):
                    t2i_image = gr.Image(
                        label=" ",
                        type="pil",
                        format="png",
                        height=375,
                    )

                    t2i_file = gr.File(
                        label=" ",
                        visible=False,
                    )

                    with gr.Row():
                        t2i_used_seed = gr.Number(
                            label=" ",
                            interactive=False,
                            visible=False,
                        )

                        t2i_timing = gr.Textbox(
                            label=" ",
                            interactive=False,
                            visible=False,
                        )

            gr.Examples(
                examples=[
                    [
                        'A neon shop sign that reads "QWEN IMAGE 2.1", '
                        "rainy night, reflections on wet pavement",
                        False,
                        "16:9 — 1376x768",
                    ],
                    [
                        "A capybara wearing a wizard hat, reading a book by "
                        "candlelight, oil painting",
                        False,
                        "1:1 — 1024x1024",
                    ],
                    [
                        "A cute cartoon dragon sticker",
                        True,
                        "1:1 — 1024x1024",
                    ],
                    [
                        "A glossy 3D icon of a paper plane, soft studio lighting",
                        True,
                        "1:1 — 1024x1024",
                    ],
                    [
                        "A handwritten chalkboard menu in a Parisian cafe, "
                        "legible French text, warm morning light",
                        False,
                        "3:4 — 896x1200",
                    ],
                ],
                inputs=[
                    t2i_prompt,
                    t2i_transparent,
                    t2i_size,
                ],
                outputs=[
                    t2i_image,
                    t2i_file,
                    t2i_used_seed,
                    t2i_timing,
                ],
                fn=generate,
                cache_examples=True,
                cache_mode="lazy",
                visible=False,
            )

        with gr.Tab(" "):
            with gr.Row():
                with gr.Column(scale=1):
                    edit_instruction = gr.Textbox(
                        label=" ",
                        lines=6,
                    )

                    edit_refs = gr.Gallery(
                        label=" ",
                        type="filepath",
                        columns=5,
                        height=344,
                        interactive=True,
                    )

                    edit_size = gr.Radio(
                        label=" ",
                        choices=list(EDIT_SIZE_PRESETS),
                        value=DEFAULT_EDIT_SIZE,
                    )

                    edit_run = gr.Button(
                        " ",
                        variant="primary",
                    )

                    with gr.Accordion(" ", open=False):
                        edit_steps = gr.Slider(
                            10,
                            50,
                            value=25,
                            step=1,
                            label=" ",
                        )

                        edit_cfg = gr.Slider(
                            1.0,
                            10.0,
                            value=1.1,
                            step=0.1,
                            label=" ",
                        )

                        edit_negative = gr.Textbox(
                            label=" ",
                            lines=2,
                        )

                        edit_seed = gr.Slider(
                            0,
                            MAX_SEED,
                            value=0,
                            step=1,
                            label=" ",
                        )

                        edit_random = gr.Checkbox(
                            label=" ",
                            value=True,
                        )

                with gr.Column(scale=1):
                    edit_image = gr.Image(
                        label=" ",
                        type="pil",
                        format="png",
                        height=344,
                    )

                    edit_file = gr.File(
                        label=" ",
                        visible=False,
                    )

                    with gr.Row():
                        edit_used_seed = gr.Number(
                            label=" ",
                            interactive=False,
                            visible=False,
                        )

                        edit_timing = gr.Textbox(
                            label=" ",
                            interactive=False,
                            visible=False,
                        )

            gr.Markdown(" ")

    t2i_run.click(
        fn=generate,
        inputs=[
            t2i_prompt,
            t2i_transparent,
            t2i_size,
            t2i_steps,
            t2i_cfg,
            t2i_negative,
            t2i_seed,
            t2i_random,
        ],
        outputs=[
            t2i_image,
            t2i_file,
            t2i_used_seed,
            t2i_timing,
        ],
        api_name="generate",
    )

    t2i_prompt.submit(
        fn=generate,
        inputs=[
            t2i_prompt,
            t2i_transparent,
            t2i_size,
            t2i_steps,
            t2i_cfg,
            t2i_negative,
            t2i_seed,
            t2i_random,
        ],
        outputs=[
            t2i_image,
            t2i_file,
            t2i_used_seed,
            t2i_timing,
        ],
        api_name=False,
    )

    edit_run.click(
        fn=edit,
        inputs=[
            edit_instruction,
            edit_refs,
            edit_size,
            edit_steps,
            edit_cfg,
            edit_negative,
            edit_seed,
            edit_random,
        ],
        outputs=[
            edit_image,
            edit_file,
            edit_used_seed,
            edit_timing,
        ],
        api_name="edit",
    )


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