"""MiniMax-H3 360° Orbit LoRA — the denoising half, inlined as a gr.Workflow bound function.

Adapted from pablodawson/minimax-h3-360-orbit/app.py: instead of serving a gr.Blocks UI, this module
exposes `orbit()` for `gr.Workflow(bind={"orbit": orbit})`. The Qwen3-VL conditioner half still runs in
`multimodalart/qwen3vl-conditioner` and is called over the gradio API — that half cannot be inlined
(62.14 GiB), and it is a separate Space, not the orbit denoiser this module runs locally on ZeroGPU.
"""

from __future__ import annotations

import os
import tempfile
import time
import traceback
from functools import cache

# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the DiT load can
# happen at startup rather than on GPU time.
import spaces
import gradio as gr
import torch

VERSION = "h3-360-orbit/workflow-1"

# The AdaLN-pruned `transformer` (fl2va) partition in diffusers layout — the same weights this LoRA was
# trained against, minus the AdaLN input matrices the adapter never touches anyway.
MODEL_REPO = os.environ.get("H3_MODEL_REPO", "multimodalart/MiniMax-H3-Pruned")
LORA_REPO = os.environ.get("H3_LORA_REPO", "pablodawson/MiniMax-H3-360-Orbit-LoRA")
LORA_FILE = os.environ.get("H3_LORA_FILE", "minimax_h3_flf2v_lora_v1.safetensors")
ADAPTER = "orbit"
# "Start with the final checkpoint at strength 1.0" — the model card's own recommendation.
DEFAULT_LORA_SCALE = float(os.environ.get("H3_LORA_SCALE", "1.0"))

CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
# `lazy` moves the partition onto the card on the first GPU call and leaves it there. Startup placement
# is not an option: `spaces`' startup `torch.pack()` would need the weights plus their pack on disk and
# that does not fit the storage quota.
PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
# cuDNN's fused attention, 10-20% faster than the SDPA default on this pool and needs nothing installed.
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
MIN_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MIN", "120"))
MAX_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MAX", "1200"))

# Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half
# does not know is rejected there and surfaces as a failure here. Only 1:1 scaffolds the LoRA's own
# training domain; the other shapes are offered and the card itself flags them as untested.
CANVASES = {
    # 16:9
    "960x544 · 16:9 fast": (544, 960),
    "1024x576 · 16:9 fast": (576, 1024),
    "1152x640 · 16:9": (640, 1152),
    "1280x704 · 16:9": (704, 1280),
    "1344x768 · 16:9 full": (768, 1344),
    # 9:16
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    # 1:1
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    "1024x1024 · 1:1 max": (1024, 1024),
    # 4:3 / 3:4
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    # 21:9
    "1152x512 · 21:9 fast": (512, 1152),
    "1536x672 · 21:9 full": (672, 1536),
}
# The LoRA's own training bucket — 768×768 square clips, the resolution every result in its card used.
DEFAULT_CANVAS = "768x768 · 1:1 full"

# The single prompt every clip in the training set was captioned with. The card is explicit: use it
# verbatim.
ORBIT_PROMPT = (
    "One frozen instant. Only the camera moves. In a continuous 360 orbit. Preserve every person and "
    "object in exactly the same world position, orientation, shape and pose throughout the shot. "
    "Airborne objects remain suspended at the captured height and angle: no wobbling, shaking, "
    "spinning, drifting, falling or continued action. Keep faces, hands, clothing, liquids and the "
    "background motionless while retaining their natural appearance. Camera parallax is the only source "
    "of apparent movement. No cuts, zoom, morphing or added objects."
)

FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
DEFAULT_STEPS = 28
DEFAULT_SECONDS = 3
MIN_UI_DURATION = 2

# Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows
# for the matmuls, quadratic for the attention. The coefficients are the fit the split MiniMax-H3
# Spaces measured on this pool.
STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.0
# The lazy `PIPE.to("cuda")` a cold worker pays inside its first GPU call.
PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "45"))
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
CANVAS_MULTIPLE = 32
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124

OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "h3-360-orbit-out")
os.makedirs(OUTPUT_DIR, exist_ok=True)

PIPE = None
LOAD_ERROR: str | None = None
LORA_STATUS: str = ""


def snap_frames(seconds: float) -> int:
    """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps."""
    frames = max(1, round(float(seconds) * FPS))
    while frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK:
        frames += 1
    return frames


def lower_duration_floor(seconds: float = MIN_UI_DURATION) -> None:
    """Let the pipeline generate below its 5 s floor. 56 frames (2.33 s) is fine on the pruned checkpoint."""
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline

    MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))


def target_rows(height: int, width: int, num_frames: int) -> int:
    """The generated rows of the packed sequence: video patched `(1, 2, 2)`, plus two audio rows per latent."""
    latents = 5 * ((num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) + 2
    video = latents * (height // CANVAS_MULTIPLE) * (width // CANVAS_MULTIPLE)
    return video + round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS


# ── LoRA loading (runtime PEFT adapter) ──────────────────────────────────────
#
# The adapter targets the original checkpoint's module names; diffusers serves the converted port, so
# every name — and, for two of them, the row layout of `lora_B` — has to be pushed through the same
# transforms the conversion script applied to the base weights:
#
#   * `blocks.` -> `transformer_blocks.`, `token_refiner.blocks.` -> `token_refiner.refiner_blocks.`,
#     `attn.out_proj` -> `attn.to_out.0`, `mlp.fc2` -> `ff.net.2` (pure renames),
#   * `mlp.fc1` -> `ff.net.0.proj` with its two fused halves swapped (diffusers' SwiGLU reads
#     `[value; gate]` where the checkpoint stores `[gate; value]`),
#   * `attn.qkv_proj` -> `to_q` / `to_k` / `to_v`: split `lora_B`'s rows into contiguous thirds. This
#     file is an ai-toolkit export (`diffusion_model.…` keys, no `.default.` infix), so it takes
#     contiguous thirds with no de-interleaving.
#
# Row transforms only ever touch `lora_B`, so the three projections share one `lora_A`.


def _lora_target_name(source_name: str) -> str:
    """Map an original-checkpoint module path to the diffusers one (no `.lora_A/B.*` suffix)."""
    if source_name.startswith("token_refiner.blocks."):
        return source_name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
    if source_name.startswith("blocks."):
        return source_name.replace("blocks.", "transformer_blocks.", 1)
    return source_name


def _lora_modules(name: str, a_weight, b_weight, inner_dim: int):
    """Yield `(diffusers_module_path, lora_A, lora_B)` for one original module path."""
    target = _lora_target_name(name)

    if target.endswith(".attn.qkv_proj"):
        prefix = target.removesuffix("qkv_proj")
        for kind, part in zip(("q", "k", "v"), b_weight.split(inner_dim, dim=0)):
            yield f"{prefix}to_{kind}", a_weight, part.contiguous()
    elif target.endswith(".mlp.fc1"):
        # SwiGLU gate/value swap: the checkpoint stores [gate, value]; diffusers stores [value, gate].
        gate, value = b_weight.chunk(2, dim=0)
        yield target.replace(".mlp.fc1", ".ff.net.0.proj"), a_weight, torch.cat([value, gate]).contiguous()
    elif target.endswith(".mlp.fc2"):
        yield target.replace(".mlp.fc2", ".ff.net.2"), a_weight, b_weight
    elif target.endswith(".attn.out_proj"):
        yield target.replace(".attn.out_proj", ".attn.to_out.0"), a_weight, b_weight
    else:
        yield target, a_weight, b_weight


def build_lora_state_dict(transformer) -> tuple[dict, int, int]:
    """Download the LoRA and remap it into a PEFT-format state dict for the diffusers transformer."""
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    raw = load_file(hf_hub_download(LORA_REPO, LORA_FILE))

    prefix, suffix_a, suffix_b = "diffusion_model.", ".lora_A.weight", ".lora_B.weight"
    unexpected = [k for k in raw if not (k.startswith(prefix) and k.endswith((suffix_a, suffix_b)))]
    if unexpected:
        raise ValueError(f"{LORA_FILE} holds {len(unexpected)} unexpected tensors, e.g. {unexpected[:5]}")
    bases = sorted({k[len(prefix) : -len(suffix_a)] for k in raw if k.endswith(suffix_a)})
    if not bases:
        raise ValueError(f"No `{prefix}*{suffix_a}` / `{suffix_b}` pairs found in {LORA_FILE}")

    ranks = set()
    for name in bases:
        if f"{prefix}{name}{suffix_b}" not in raw:
            raise ValueError(f"LoRA is missing the lora_B twin of {prefix}{name}{suffix_a}")
        ranks.add(raw[f"{prefix}{name}{suffix_a}"].shape[0])
    if len(ranks) != 1:
        raise ValueError(f"LoRA mixes ranks {sorted(ranks)}; this loader assumes a single rank")
    rank = ranks.pop()

    inner_dim = transformer.config.num_attention_heads * transformer.config.attention_head_dim
    base_shapes = {key: tuple(value.shape) for key, value in transformer.state_dict().items()}

    state_dict: dict[str, torch.Tensor] = {}
    missed: list[str] = []
    for name in bases:
        a_weight = raw[f"{prefix}{name}{suffix_a}"]
        b_weight = raw[f"{prefix}{name}{suffix_b}"]
        for module, a_part, b_part in _lora_modules(name, a_weight, b_weight, inner_dim):
            base = base_shapes.get(f"{module}.weight")
            if base is None:
                missed.append(module)
                continue
            if (b_part.shape[0], a_part.shape[1]) != base:
                raise ValueError(
                    f"LoRA delta for `{module}` would be {(b_part.shape[0], a_part.shape[1])}, "
                    f"base weight is {base}"
                )
            state_dict[f"{module}.lora_A.weight"] = a_part
            state_dict[f"{module}.lora_B.weight"] = b_part
    if missed:
        raise ValueError(
            f"{len(missed)} LoRA targets matched no transformer weight, e.g. {missed[:5]}. "
            "The LoRA and the diffusers transformer disagree on module naming."
        )

    return state_dict, rank, len(bases)


def load_and_apply_lora(transformer) -> str:
    """Attach the LoRA as a PEFT adapter so its strength stays a per-request knob."""
    state_dict, rank, targets = build_lora_state_dict(transformer)
    transformer.load_lora_adapter(state_dict, prefix=None, adapter_name=ADAPTER)
    transformer.set_adapters([ADAPTER], [DEFAULT_LORA_SCALE])
    return (
        f"LoRA attached · {targets} checkpoint targets -> {len(state_dict) // 2} diffusers modules, "
        f"rank {rank}, alpha == rank · default strength {DEFAULT_LORA_SCALE:g} · {LORA_FILE}"
    )


def load_models() -> str | None:
    """Load the denoising half at startup, but *not* onto the card."""
    global PIPE, LOAD_ERROR, LORA_STATUS

    if PIPE is not None or LOAD_ERROR is not None:
        return LOAD_ERROR

    started = time.time()
    try:
        from diffusers import ComponentsManager

        from h3_split_blocks import MiniMaxH3GeneratorBlocks

        lower_duration_floor()
        manager = ComponentsManager()
        blocks = MiniMaxH3GeneratorBlocks()
        print(f"[{VERSION}] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
        pipe.load_components(dtype=torch.bfloat16, trust_remote_code=True)

        # Both VAEs first, and explicitly, because `set_attention_backend` also sets the registry's
        # *global* backend: the float32 audio VAE has no cuDNN kernel and would die in its decoder.
        pipe.vae.set_attention_backend("native")
        pipe.audio_vae.set_attention_backend("native")
        pipe.transformer.set_attention_backend(ATTENTION)

        # Before any GPU placement, so the adapter's parameters travel with the transformer.
        LORA_STATUS = load_and_apply_lora(pipe.transformer)
        print(f"[{VERSION}] {LORA_STATUS}", flush=True)

        PIPE = pipe
        print(f"[{VERSION}] ready in {time.time() - started:.0f}s", flush=True)
    except Exception as error:
        traceback.print_exc()
        LOAD_ERROR = (
            f"Loading `{MODEL_REPO}` + `{LORA_REPO}` failed after {time.time() - started:.0f}s: "
            f"{type(error).__name__}: {error}"
        )
    return LOAD_ERROR


@cache
def conditioner():
    """The other half, over the gradio API. `gradio_client` attaches the caller's own ZeroGPU token per
    call, so the conditioner's booking is billed to whoever asked for the video."""
    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def _caller_ip_token() -> str | None:
    try:
        from gradio.context import LocalContext

        request = LocalContext.request.get()
        return request.headers.get("x-ip-token") if request is not None else None
    except Exception:
        return None


def encode_remote(prompt, image_path, last_image_path, canvas, num_frames, ip_token=None):
    """`/encode` on the conditioner Space: a safetensors file holding `prompt_embeds` + `text_token_tags`,
    with the resolved `height` / `width` / `num_frames` in its metadata, plus the plan."""
    from gradio_client import Client, handle_file
    from safetensors import safe_open

    def client():
        if not ip_token:
            return conditioner()
        return Client(CONDITIONER_SPACE, headers={"x-ip-token": ip_token})

    def call():
        return client().predict(
            prompt=prompt,
            image_path=handle_file(image_path) if image_path else None,
            last_image_path=handle_file(last_image_path) if last_image_path else None,
            canvas=canvas,
            num_frames=num_frames,
            rewrite_prompt=False,
            api_name="/encode",
        )

    try:
        path, plan = call()
    except Exception as first:
        print(f"[{VERSION}] retrying the conditioner with a fresh client after: {first}", flush=True)
        conditioner.cache_clear()
        path, plan = call()

    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


def _as_path(value):
    """Unwrap whatever the workflow executor hands a media port: path string, {path}, {url} or FileValue."""
    if isinstance(value, dict):
        value = value.get("path") or value.get("url") or ""
        value = value.removeprefix("/gradio_api/file=")
    return value or None


def _localize(path: str) -> str:
    """Download remote image URLs to a local file; the pipeline and the cover-crop need a real path."""
    if not path.startswith(("http://", "https://")):
        return path
    import requests

    response = requests.get(path, timeout=120)
    response.raise_for_status()
    suffix = os.path.splitext(path.split("?")[0])[1] or ".jpg"
    fd, tmp = tempfile.mkstemp(suffix=suffix, dir=OUTPUT_DIR)
    with os.fdopen(fd, "wb") as handle:
        handle.write(response.content)
    return tmp


def _cover_crop(image_path: str, canvas_label: str) -> str:
    """Center cover-crop the file at `image_path` to `canvas_label`'s aspect ratio, in place."""
    from PIL import Image as _Image

    h, w = CANVASES[canvas_label]
    target = w / h
    img = _Image.open(image_path)
    if abs(img.width / img.height - target) <= 1e-3:
        return image_path
    if img.width / img.height > target:
        new_w = int(img.height * target)
        left = (img.width - new_w) // 2
        img = img.crop((left, 0, left + new_w, img.height))
    else:
        new_h = int(img.width / target)
        top = (img.height - new_h) // 2
        img = img.crop((0, top, img.width, top + new_h))
    img.save(image_path)
    return image_path


def get_duration(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, lora_scale, *_, **__):
    """Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates."""
    height, width, num_frames, steps = int(height), int(width), int(num_frames), int(steps)
    patches = (height // CANVAS_MULTIPLE) * (width // CANVAS_MULTIPLE)
    rows = target_rows(height, width, num_frames) + 2 * patches
    denoise = int(steps) * (STEP_LINEAR * rows + STEP_QUADRATIC * rows**2) * SAFETY
    decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
    total = PLACEMENT_ALLOWANCE + denoise + decode + 10
    duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
    print(f"[{VERSION}] S={rows} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
    return duration


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, steps, seed, lora_scale):
    """The only thing on GPU time: the keyframe encodes, the packed-sequence denoise loop and the decoders."""
    if PLACEMENT == "lazy":
        PIPE.to("cuda")

    # Per request, inside the worker: the adapter's scale is a python attribute on its PEFT layers, so
    # setting it on this fork's copy leaves every concurrent request with its own strength.
    PIPE.transformer.set_adapters([ADAPTER], [float(lora_scale)])

    state = PIPE(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        image=image,
        last_image=last_image,
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=int(steps),
        generator=torch.Generator("cpu").manual_seed(int(seed)),
    )
    return state.get("videos")[0]


def orbit(image=None, canvas=DEFAULT_CANVAS, steps=DEFAULT_STEPS, seed=904231, lora_scale=DEFAULT_LORA_SCALE, duration=DEFAULT_SECONDS):
    """Orbit the camera 360° around one photo, with the scene frozen and the first frame exact.

    Bound into the workflow as the `orbit` operator node. The photo is fed to MiniMax-H3 as **both**
    the first and the last keyframe — the wiring the LoRA was trained with and what closes the orbit
    on its own first frame.

    Returns:
        `(video_path, report)` — path to the generated MP4 (video only; the card recommends audio
        off) and a report line.
    """
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading — watch the Space logs and retry shortly.")

    from PIL import Image, ImageOps

    from diffusers.utils import encode_video

    path = _as_path(image)
    if not path:
        raise gr.Error("Upload a photo to orbit around first.")
    if not canvas or canvas not in CANVASES:
        raise gr.Error(f"Pick a supported canvas. {canvas!r} is not one.")

    # Workflow media ports can arrive as remote URLs (e.g. the preset example image); the cover-crop
    # and the conditioner upload both need a local file.
    first = _cover_crop(_localize(path), canvas)

    num_frames = snap_frames(float(duration))

    conditioned = time.time()
    try:
        prompt_embeds, text_token_tags, metadata, plan = encode_remote(
            ORBIT_PROMPT, first, first, canvas, num_frames, ip_token=_caller_ip_token()
        )
    except gr.Error:
        raise
    except Exception as error:
        traceback.print_exc()
        raise gr.Error(
            f"The conditioner ({CONDITIONER_SPACE}) failed with `{type(error).__name__}: {error}`. "
            "Its logs carry the full traceback."
        ) from error
    condition_seconds = time.time() - conditioned
    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))

    def keyframe(p):
        return ImageOps.exif_transpose(Image.open(p)).convert("RGB")

    started = time.time()
    frames = _generate(
        prompt_embeds,
        text_token_tags,
        keyframe(first),
        keyframe(first),  # the same image, again: the last keyframe is what closes the orbit
        height,
        width,
        num_frames,
        int(float(steps)),
        int(float(seed)),
        float(lora_scale),
    )
    generate_seconds = time.time() - started

    # The card recommends audio off; the joint audio latents are still generated and decoded by the
    # pipeline, but nothing here hands them to the mux.
    out = os.path.join(OUTPUT_DIR, f"h3-360-orbit-{int(time.time() * 1000)}.mp4")
    encode_video(frames, fps=FPS, output_path=out)

    report = (
        f"{width}x{height} · {num_frames} frames ({num_frames / FPS:.3f} s) · {int(float(steps))} steps · "
        f"strength {float(lora_scale):g} · conditioner {condition_seconds:.0f}s "
        f"({plan['num_text_tokens']} tokens) · denoise + decode {generate_seconds:.0f}s "
        f"({generate_seconds / max(1, int(float(steps))):.1f} s/step) · seed {int(float(seed))}"
    )
    print(f"[{VERSION}] {report}", flush=True)
    # The workflow canvas only renders a media fn output that carries a servable `url` — a bare local
    # path falls through to its text branch and the video subject shows nothing. OUTPUT_DIR is under
    # the temp dir, which Workflow.launch auto-adds to allowed_paths (and app.py adds it explicitly).
    import urllib.parse

    video = {
        "url": f"/gradio_api/file={urllib.parse.quote(out)}",
        "orig_name": os.path.basename(out),
        "mime_type": "video/mp4",
    }
    return video, report


load_models()