"""P-H3-Image on ZeroGPU (`xlarge`: a full RTX PRO 6000 Blackwell, 96 GB).

The pipeline is loaded once at startup, one component at a time (the host has ~97 GB), and slimmed to fit the slice
with p_h3_image_fast.py (kernels only, no torch.compile): AdaLN tables instead of the transformer's AdaLN projections,
FP8 linears everywhere, fused Triton DiT blocks, batched VAE tiles, and for the prompt upsampler a weight-only FP8 GEMV,
decode attention and speculative decoding with its MTP head in one CUDA graph per step.
"""

import json
import os
import random
import tempfile
import time

os.environ.setdefault("DIFFUSERS_TRUST_REMOTE_KERNELS", "true")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import gradio as gr
import spaces
import torch
from huggingface_hub import snapshot_download
from PIL import Image, ImageOps
from torchvision.io import encode_jpeg

import p_h3_image_fast as fast
from diffusers import ModularPipeline
from examples import EXAMPLES
from theme import CSS, HEAD, THEME, render_footer, render_header

REPO, REVISION = "PrunaAI/P-H3-Image", "ee4a685de1fbd4337a50bcf9b3c407e4f8463d49"
# Canvases (width, height), multiples of 32: 4 and 8 MP as tested, 16 MP from the same ratios.
RATIOS = {"1:1": 1.0, "4:3": 4 / 3, "3:4": 3 / 4, "3:2": 3 / 2, "2:3": 2 / 3, "16:9": 16 / 9, "9:16": 9 / 16}
T2I_RESOLUTIONS, EDIT_RESOLUTIONS = ["4 MP", "8 MP", "16 MP"], ["4 MP", "8 MP"]


def _canvas(ratio: float, megapixels: int) -> tuple[int, int]:
    pixels = megapixels * 1024**2
    return round((pixels * ratio) ** 0.5 / 32) * 32, round((pixels / ratio) ** 0.5 / 32) * 32


SIZES = {
    "4 MP": {"1:1": (2048, 2048), "4:3": (2368, 1760), "3:4": (1760, 2368), "3:2": (2496, 1664), "2:3": (1664, 2496),
             "16:9": (2720, 1536), "9:16": (1536, 2720)},
    "8 MP": {"1:1": (2912, 2912), "4:3": (3360, 2496), "3:4": (2496, 3360), "3:2": (3552, 2368), "2:3": (2368, 3552),
             "16:9": (3872, 2176), "9:16": (2176, 3872)},
    "16 MP": {name: _canvas(ratio, 16) for name, ratio in RATIOS.items()},
}
MAX_IMAGES = 8
# All references together stay within 8 MP: one image keeps up to 8 MP, eight get 1 MP each.
REFERENCE_BUDGET = 8 * 1024**2


def load_pipeline():
    path = snapshot_download(REPO, revision=REVISION, max_workers=16)
    with open(os.path.join(os.path.dirname(__file__), "adaln_calls.json")) as handle:
        calls = json.load(handle)
    pipe, mtp = fast.load_slim_pipeline(path, calls, log=lambda message: print(message, flush=True))
    # The same components without the prompt upsampler step, for requests that skip upsampling.
    blocks = ModularPipeline.from_pretrained(path).blocks
    blocks.sub_blocks.pop("prompt_upsampler")
    plain = blocks.init_pipeline(path)
    plain.update_components(**{name: getattr(pipe, name) for name in plain.component_names
                               if getattr(pipe, name, None) is not None})
    return pipe, plain, mtp


PIPE, PLAIN, MTP = load_pipeline()
UPSAMPLER_GENERATE = PIPE.prompt_upsampler.generate
STATE = {"kernels": False, "decoder": None, "owner": None}


def _prepare_gpu():
    """Per GPU call. Once per worker process: the upsampler's hub kernels, and the speculative decoder with its CUDA
    graph (captured on the first request, then replayed; rebuilt if the process or device changes)."""
    upsampler = PIPE.prompt_upsampler
    if not STATE["kernels"]:
        from kernels import Mode

        upsampler.set_use_kernels(True, mode=Mode.INFERENCE)
        STATE["kernels"] = True
    # The graph bakes in weight addresses: rebuild if ZeroGPU placed the weights elsewhere since the capture.
    weights = (upsampler.lm_head.weight, MTP.fc.weight, upsampler.model.language_model.embed_tokens.weight,
               upsampler.model.language_model.layers[0].mlp.down_proj.weight)
    owner = (os.getpid(), torch.cuda.current_device(), tuple(w.data_ptr() for w in weights))
    if STATE["decoder"] is None or STATE["owner"] != owner:
        upsampler.generate = UPSAMPLER_GENERATE
        STATE["decoder"], STATE["owner"] = fast.MTPGreedyGenerate(upsampler, MTP), owner
    decoder = STATE["decoder"]
    timing = {}

    def generate(*args, **kwargs):
        torch.cuda.synchronize()
        begin = time.perf_counter()
        out = decoder.generate(*args, **kwargs)
        torch.cuda.synchronize()
        timing["upsampler"] = time.perf_counter() - begin
        timing["tokens"] = decoder.stats.get("tokens", 0)
        return out

    upsampler.generate = generate
    return timing


def _duration(mode, prompt, images, aspect, resolution, upsample, seed):
    # Measured on a full card: ~6 / 12 / 25 s for 4 / 8 / 16 MP text-to-image, edits a bit more per image. Shorter
    # requests queue faster on ZeroGPU.
    base = {"4 MP": 25, "8 MP": 35, "16 MP": 60}[resolution]
    return base + (10 + 5 * len(images) if mode == "edit" else 0)


@spaces.GPU(size="xlarge", duration=_duration)
def _render(mode, prompt, images, aspect, resolution, upsample, seed):
    timing = _prepare_gpu()
    torch.cuda.synchronize()
    start = time.perf_counter()
    kwargs = {"prompt": prompt, "generator": torch.Generator().manual_seed(seed)}
    if mode == "edit":
        # The first image's aspect ratio at the chosen size.
        width, height = images[0].size
        kwargs["image"] = images
        kwargs["width"], kwargs["height"] = _canvas(width / height, int(resolution.split()[0]))
    else:
        kwargs["width"], kwargs["height"] = SIZES[resolution][aspect]
    result = (PIPE if upsample else PLAIN)(**kwargs, output_type="pt", output=["images", "prompt"])
    path, size = _to_jpeg(result["images"])
    total = time.perf_counter() - start
    return path, size, result["prompt"], {"total": total, **timing}


def _to_jpeg(images: torch.Tensor) -> tuple[str, tuple[int, int]]:
    """The decoded image as a JPEG file: uint8 and nvjpeg on the GPU, so only the encoded bytes leave the worker."""
    pixels = (images[0].clamp(0, 1) * 255).round().to(torch.uint8).contiguous()
    try:
        data = encode_jpeg(pixels, quality=93)
    except Exception:  # noqa: BLE001 - no nvjpeg: encode on the CPU
        data = encode_jpeg(pixels.cpu(), quality=93)
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as handle:
        handle.write(data.cpu().numpy().tobytes())
    return handle.name, (pixels.shape[2], pixels.shape[1])


@spaces.GPU(size="xlarge", duration=300)
def _warmup():
    """Triton compiles and autotunes, and the hub kernels load, before the first visitor."""
    image = Image.new("RGB", (1024, 1024), (120, 140, 160))
    timing = _prepare_gpu()
    PIPE(prompt="a lighthouse on a cliff at dusk", width=1024, height=1024, generator=torch.Generator().manual_seed(0),
         output="images")
    PIPE(prompt="make it night", image=[image], width=1024, height=1024, generator=torch.Generator().manual_seed(0),
         output="images")
    return timing


try:
    _warmup()
except Exception as error:  # noqa: BLE001 - the demo still works, the first request just pays the warmup
    print(f"warmup failed: {error!r}", flush=True)


def _info(timing, size, seed):
    upsample, tokens = timing.get("upsampler", 0.0), timing.get("tokens", 0)
    rate = f", {tokens} tokens at {tokens / upsample:.0f} tokens/s" if upsample and tokens else ""
    head = f"{size[0]}×{size[1]}, seed {seed} · {timing['total']:.1f} s on the GPU"
    if not upsample:
        return f"{head} (no prompt upsampling)"
    return f"{head} ({upsample:.1f} s prompt upsampling{rate}, {timing['total'] - upsample:.1f} s rendering)"


def text_to_image(prompt, aspect, resolution, upsample):
    if not prompt or not prompt.strip():
        raise gr.Error("Write a prompt first.")
    seed = random.randint(0, 2**31 - 1)
    try:
        path, size, upsampled, timing = _render("t2i", prompt.strip(), [], aspect, resolution, upsample, seed)
    except RuntimeError as error:
        raise gr.Error(str(error).split("\n")[0])
    return path, upsampled, _info(timing, size, seed)


def _fit(image: Image.Image, pixels: float) -> Image.Image:
    """Downscale to at most `pixels` (the pipeline keeps references at their own size, from 1 MP up)."""
    scale = (pixels / (image.width * image.height)) ** 0.5
    if scale >= 1:
        return image
    return image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale))), Image.LANCZOS)


def edit(gallery, prompt, resolution, upsample):
    images = []
    for item in gallery or []:
        item = item[0] if isinstance(item, (tuple, list)) else item
        item = Image.open(item) if isinstance(item, str) else item
        images.append(ImageOps.exif_transpose(item).convert("RGB"))
    if not images:
        raise gr.Error("Add at least one image; the first one is edited.")
    if len(images) > MAX_IMAGES:
        raise gr.Error(f"At most {MAX_IMAGES} images.")
    if not prompt or not prompt.strip():
        raise gr.Error("Describe the edit first.")
    images = [_fit(image, REFERENCE_BUDGET / len(images)) for image in images]
    seed = random.randint(0, 2**31 - 1)
    try:
        path, size, upsampled, timing = _render("edit", prompt.strip(), images, None, resolution, upsample, seed)
    except (RuntimeError, ValueError) as error:
        raise gr.Error(str(error).split("\n")[0])
    return path, upsampled, _info(timing, size, seed)


with gr.Blocks(title="P-H3-Image", fill_width=True) as demo:
    render_header()
    with gr.Tabs(elem_classes="main-tabs"):
        with gr.Tab("Text to image"):
            with gr.Row():
                with gr.Column():
                    t_prompt = gr.Textbox(label="Prompt", lines=4)
                    with gr.Row():
                        t_aspect = gr.Dropdown(list(RATIOS), value="4:3", label="Aspect ratio", filterable=False)
                        t_resolution = gr.Radio(T2I_RESOLUTIONS, value="8 MP", label="Resolution")
                    t_upsample = gr.Checkbox(value=True, label="Prompt upsampling")
                    t_button = gr.Button("Generate", variant="primary")
                with gr.Column():
                    t_image = gr.Image(label="Result", type="filepath")
                    t_info = gr.Markdown()
                    with gr.Accordion("Upsampled prompt", open=False):
                        t_upsampled = gr.Textbox(show_label=False, lines=10)
            gr.Examples([[prompt, aspect] for _, aspect, prompt in EXAMPLES], inputs=[t_prompt, t_aspect],
                        example_labels=[label for label, _, _ in EXAMPLES])
            t_button.click(text_to_image, [t_prompt, t_aspect, t_resolution, t_upsample], [t_image, t_upsampled, t_info])
        with gr.Tab("Edit"):
            with gr.Row():
                with gr.Column():
                    e_images = gr.Gallery(label=f"Images (the first one is edited, up to {MAX_IMAGES})", type="pil",
                                          columns=4, height=260, interactive=True)
                    e_prompt = gr.Textbox(label="Edit", lines=3, placeholder="e.g. put the fox from <Picture 2> on the bench")
                    e_resolution = gr.Radio(EDIT_RESOLUTIONS, value="8 MP", label="Resolution",
                                            info="The first image's aspect ratio at this size.")
                    e_upsample = gr.Checkbox(value=True, label="Prompt upsampling")
                    e_button = gr.Button("Edit", variant="primary")
                with gr.Column():
                    e_image = gr.Image(label="Result", type="filepath")
                    e_info = gr.Markdown()
                    with gr.Accordion("Upsampled prompt", open=False):
                        e_upsampled = gr.Textbox(show_label=False, lines=10)
            e_button.click(edit, [e_images, e_prompt, e_resolution, e_upsample], [e_image, e_upsampled, e_info])
    render_footer()

demo.queue(max_size=20).launch(theme=THEME, css=CSS, head=HEAD)