"""P-H3-Image Studio — gr.Server studio + gr.Workflow canvas.

The bespoke studio at / replaces the classic form UI; the backend is unchanged from
PrunaAI/P-H3-Image: 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. The studio's
Generate/Edit calls and the workflow canvas at /workflow all run through the same _render.

Server + workflow structure adapted from https://huggingface.co/spaces/briaai/Fibo-2.0-editing.
"""

import base64
import json
import os
import random
import shutil
import tempfile
import time
from urllib.parse import unquote, urlparse

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
import uvicorn
from fastapi.responses import FileResponse, HTMLResponse
from gradio.utils import get_upload_folder
from gradio_client import utils as client_utils
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

REPO, REVISION = "PrunaAI/P-H3-Image", "ee4a685de1fbd4337a50bcf9b3c407e4f8463d49"
APP_DIR = os.path.dirname(os.path.abspath(__file__))
# 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"]
MAX_SEED = 2**31 - 1


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(APP_DIR, "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 _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)


# --------------------------------------------------------------------------- helpers


def _resolve_seed(seed, randomize) -> int:
    return random.randint(0, MAX_SEED) if randomize else min(max(int(float(seed)), 0), MAX_SEED)


def _check_choice(value, choices, name):
    if value not in choices:
        raise gr.Error(f"{name} must be one of: {', '.join(choices)}.")
    return value


def _as_data_uri(path: str) -> str:
    """Inline a result image so the studio page never needs a file route."""
    with open(path, "rb") as handle:
        return "data:image/jpeg;base64," + base64.b64encode(handle.read()).decode()


def _requested_path(value) -> str:
    """Turn whatever the canvas/API client sends into a path."""
    if isinstance(value, str):
        return value
    if not isinstance(value, dict):
        raise gr.Error(f"Expected an image, got {type(value).__name__}")
    if path := value.get("path"):
        return path
    url = value.get("url") or ""
    if "/gradio_api/file=" in url:
        return unquote(url.split("/gradio_api/file=", 1)[1])
    if url.startswith(("http://", "https://")):
        raise gr.Error("Upload the image instead of passing a URL.")
    return unquote(urlparse(url).path)


def _local_path(value) -> str:
    """A path this process may open: only files inside gradio's upload folder,
    where /gradio_api/upload puts uploads and _publish puts results."""
    path = os.path.realpath(_requested_path(value))
    folder = os.path.realpath(get_upload_folder())
    if os.path.commonpath([path, folder]) != folder:
        raise gr.Error("Upload the image first.")
    return path


def _open_rgb(value) -> Image.Image:
    return ImageOps.exif_transpose(Image.open(_local_path(value))).convert("RGB")


def _publish(path: str) -> str:
    """Copy a result into gradio's upload folder, which it already serves, so the workflow canvas can render it."""
    folder = get_upload_folder()
    os.makedirs(folder, exist_ok=True)
    out = os.path.join(folder, f"ph3_{os.urandom(8).hex()}.jpg")
    shutil.copyfile(path, out)
    return out


def _as_media(path: str) -> dict:
    """The dict shape the workflow canvas renders media from."""
    return {
        "path": path,
        "url": f"/gradio_api/file={client_utils.encode_file_path(path)}",
        "is_file": True,
    }


def _run(mode, prompt, images, aspect, resolution, upsample, seed):
    """One _render call with the classic app's error surface: messages, not stack traces."""
    try:
        return _render(mode, prompt, images, aspect, resolution, upsample, seed)
    except (RuntimeError, ValueError) as error:
        raise gr.Error(str(error).split("\n")[0])


# ------------------------------------------------------------------ workflow fns


def ph3_generate(prompt, aspect: str = "4:3", resolution: str = "8 MP", seed: str = "42") -> dict:
    """Text-to-image with P-H3-Image (prompt upsampling on)."""
    if not prompt or not str(prompt).strip():
        raise gr.Error("Write a prompt first.")
    _check_choice(aspect, RATIOS, "Aspect ratio")
    _check_choice(resolution, T2I_RESOLUTIONS, "Resolution")
    seed_value = min(max(int(float(seed)), 0), MAX_SEED)
    path, _, _, _ = _run("t2i", str(prompt).strip(), [], aspect, resolution, True, seed_value)
    return _as_media(_publish(path))


def ph3_edit(image, prompt, seed: str = "42") -> dict:
    """Edit an image with P-H3-Image (prompt upsampling on, 8 MP at the image's aspect ratio)."""
    if not prompt or not str(prompt).strip():
        raise gr.Error("Describe the edit first.")
    source = _fit(_open_rgb(image), REFERENCE_BUDGET)
    seed_value = min(max(int(float(seed)), 0), MAX_SEED)
    path, _, _, _ = _run("edit", str(prompt).strip(), [source], None, "8 MP", True, seed_value)
    return _as_media(_publish(path))


# ------------------------------------------------------------------ studio server

server = gr.Server(title="P-H3-Image Studio")


@server.get("/", response_class=HTMLResponse)
def _studio_page():
    with open(os.path.join(APP_DIR, "studio.html"), encoding="utf-8") as handle:
        return handle.read()


@server.api(name="generate", concurrency_limit=1)
def api_generate(
    prompt: str,
    aspect: str = "4:3",
    resolution: str = "8 MP",
    upsample: bool = True,
    seed: float = 42,
    randomize: bool = True,
) -> dict:
    """Text-to-image with P-H3-Image."""
    if not prompt or not prompt.strip():
        raise gr.Error("Write a prompt first.")
    _check_choice(aspect, RATIOS, "Aspect ratio")
    _check_choice(resolution, T2I_RESOLUTIONS, "Resolution")
    seed_value = _resolve_seed(seed, randomize)
    path, size, upsampled, timing = _run("t2i", prompt.strip(), [], aspect, resolution, bool(upsample), seed_value)
    return {"image_data": _as_data_uri(path), "seed": seed_value, "prompt": upsampled,
            "info": _info(timing, size, seed_value)}


@server.api(name="edit", concurrency_limit=1)
def api_edit(
    source: str,
    instruction: str,
    ref2: str | None = None,
    ref3: str | None = None,
    ref4: str | None = None,
    ref5: str | None = None,
    ref6: str | None = None,
    ref7: str | None = None,
    ref8: str | None = None,
    resolution: str = "8 MP",
    upsample: bool = True,
    seed: float = 42,
    randomize: bool = True,
) -> dict:
    """Edit the source (+ up to 7 extra references) with P-H3-Image."""
    if source is None:
        raise gr.Error("Upload a source image first.")
    if not instruction or not instruction.strip():
        raise gr.Error("Describe the edit first.")
    _check_choice(resolution, EDIT_RESOLUTIONS, "Resolution")
    images = [_open_rgb(path) for path in (source, ref2, ref3, ref4, ref5, ref6, ref7, ref8) if path]
    if len(images) > MAX_IMAGES:
        raise gr.Error(f"At most {MAX_IMAGES} images.")
    images = [_fit(image, REFERENCE_BUDGET / len(images)) for image in images]
    seed_value = _resolve_seed(seed, randomize)
    path, size, upsampled, timing = _run("edit", instruction.strip(), images, None, resolution, bool(upsample),
                                         seed_value)
    return {"image_data": _as_data_uri(path), "seed": seed_value, "prompt": upsampled,
            "info": _info(timing, size, seed_value)}


# The @server.api endpoints are deferred until Server.launch(), which also fires
# the ZeroGPU registration hook (pack weights + startup report) through the
# patched Blocks.launch path. Launch on a loopback port for registration, then
# serve the same app via uvicorn.
_port_env = os.environ.get("GRADIO_SERVER_PORT")
os.environ["GRADIO_SERVER_PORT"] = "7861"
try:
    server.launch(
        prevent_thread_lock=True,
        server_name="127.0.0.1",
        server_port=7861,
        quiet=True,
        show_error=False,
        ssr_mode=False,
    )
finally:
    if _port_env is None:
        del os.environ["GRADIO_SERVER_PORT"]
    else:
        os.environ["GRADIO_SERVER_PORT"] = _port_env

try:
    server.close()
except Exception:
    pass

# ------------------------------------------------------------------ workflow canvas
# Mounted only after the registration launch above. Mounting first starts the
# workflow's queue on the loopback server's event loop, and close() then strands
# it: jobs are accepted on port 7860 but never run.

workflow_demo = gr.Workflow(
    graph="workflow.json",
    bind={"ph3_generate": ph3_generate, "ph3_edit": ph3_edit},
)
workflow_demo.title = "P-H3-Image Workflow"

server = gr.mount_gradio_app(server, workflow_demo, path="/workflow", ssr_mode=False)

if __name__ == "__main__":
    uvicorn.run(server, host="0.0.0.0", port=7860)