"""MiniMax-H3 `t2va` / `fl2va`, split deployment — the denoising half."""

from __future__ import annotations

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

# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda`
# so the 72 GiB load can happen at startup rather than on GPU time.
import spaces
from fastapi.responses import HTMLResponse
from gradio import Request, Server
from gradio.data_classes import FileData


MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER_SPACE = os.environ.get(
    "H3_CONDITIONER",
    "multimodalart/qwen3vl-conditioner",
)

# `pack` places the transformer at startup, `lazy` moves everything on the first
# GPU call, `offload` hands placement to ComponentsManager.enable_auto_cpu_offload.
PLACEMENT = os.environ.get("H3_PLACEMENT", "pack").lower()

# cuDNN's fused attention is 10-20% faster than the SDPA default on this pool.
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()

GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")


# Must stay identical to the conditioner's table.
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),

    # 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),
}

DEFAULT_CANVAS = "960x544 · 16:9 fast"

FPS = 24
FRAMES_PER_CHUNK = 17
LATENTS_PER_CHUNK = 5

MIN_UI_DURATION = 2
MAX_UI_DURATION = 14


def snap_frames(seconds: float) -> int:
    """The frame count MiniMax-H3's video VAE can decode."""
    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:
    """Allow the pipeline to generate below its normal 5-second floor."""
    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import (
        MiniMaxH3ModularPipeline,
    )

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


# ----------------------------------------------------------------------
# OUTPUT DIRECTORIES
# ----------------------------------------------------------------------

# Temporary/public output used by Gradio to return the generated video.
OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "h3-outputs")

# Persistent private bucket mounted in the Space.
#
# IMPORTANT:
# Do NOT add this directory to Gradio's `allowed_paths`.
# That prevents the private archive from being intentionally exposed through
# Gradio's file-serving route.
PRIVATE_OUTPUT_DIR = "/data/private-generations"


# ----------------------------------------------------------------------
# GLOBAL STATE
# ----------------------------------------------------------------------

PIPE = None
MANAGER = None

LOAD_ERROR: str | None = None
LOADED_IN: float | None = None
LORA_STATUS: str | None = None


# ----------------------------------------------------------------------
# STATUS
# ----------------------------------------------------------------------

def status() -> str:
    if LOAD_ERROR:
        return LOAD_ERROR

    if PIPE is None:
        return (
            f"Loading `{MODEL_REPO}` "
            "(transformer + VAEs, 77.3 GB). Watch the Space logs."
        )

    import h3_aoti

    return (
        f"Ready · transformer + VAEs **bfloat16, unquantized** · "
        f"placement `{PLACEMENT}` · "
        f"attention `{ATTENTION}` · "
        f"{h3_aoti.status()} · "
        f"{LORA_STATUS or 'no LoRA'} · "
        f"loaded in {LOADED_IN:.0f}s · "
        f"conditioner `{CONDITIONER_SPACE}`"
    )


# ----------------------------------------------------------------------
# MODEL LOADING
# ----------------------------------------------------------------------

def load_models() -> str | None:
    """Load the denoising half at startup."""

    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN, LORA_STATUS

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

    started = time.time()

    try:
        import torch
        from diffusers import ComponentsManager

        from h3_split_blocks import MiniMaxH3GeneratorBlocks

        lower_duration_floor()

        manager = ComponentsManager()
        blocks = MiniMaxH3GeneratorBlocks()

        print(
            f"[gen] loading "
            f"{[c.name for c in blocks.expected_components]} "
            f"from {MODEL_REPO} ...",
            flush=True,
        )

        pipe = blocks.init_pipeline(
            MODEL_REPO,
            components_manager=manager,
            collection="h3",
        )

        pipe.load_components(dtype=torch.bfloat16)

        # Fold the Turbo LoRA into the bf16 weights before AoTI packages
        # the blocks.
        import h3_lora

        LORA_STATUS = h3_lora.apply_lora(pipe.transformer)

        if LORA_STATUS:
            print(
                f"[gen] {LORA_STATUS}",
                flush=True,
            )

        pipe.transformer.set_attention_backend(ATTENTION)

        # AoTI package.
        import h3_aoti

        h3_aoti.maybe_load(pipe.transformer)

        if PLACEMENT == "pack":
            # Only pack the transformer at startup.
            pipe.transformer.to("cuda")

        if PLACEMENT == "offload":
            manager.enable_auto_cpu_offload(device="cuda")
            _arm_decode_hooks(pipe)

        PIPE = pipe
        MANAGER = manager
        LOADED_IN = time.time() - started

        print(
            f"[gen] ready in {LOADED_IN:.0f}s",
            flush=True,
        )

    except Exception as error:
        traceback.print_exc()

        LOAD_ERROR = (
            f"**Loading `{MODEL_REPO}` failed** "
            f"after {time.time() - started:.0f}s: "
            f"`{type(error).__name__}: {error}`"
        )

    return LOAD_ERROR


# ----------------------------------------------------------------------
# OFFLOAD HOOKS
# ----------------------------------------------------------------------

def _arm_decode_hooks(pipe):
    """Make the offload hooks fire for the two VAEs."""

    for name in ("vae", "audio_vae"):
        module = getattr(pipe, name)
        inner = module.decode

        def armed(
            *args,
            _module=module,
            _decode=inner,
            **kwargs,
        ):
            hook = getattr(_module, "_hf_hook", None)

            if hook is not None:
                hook.pre_forward(_module)

            return _decode(*args, **kwargs)

        module.decode = armed


# ----------------------------------------------------------------------
# CONDITIONER
# ----------------------------------------------------------------------

@cache
def conditioner():
    """Fallback conditioner client."""

    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def conditioner_client(ip_token):
    """Create a conditioner client billed to the caller when possible."""

    if not ip_token:
        return conditioner()

    from gradio_client import Client

    return Client(
        CONDITIONER_SPACE,
        headers={"x-ip-token": ip_token},
    )


def encode_remote(
    prompt,
    image_path,
    last_image_path,
    canvas,
    num_frames,
    rewrite_prompt=False,
    ip_token=None,
):
    """Encode prompt/keyframes through the conditioner Space."""

    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = conditioner_client(ip_token).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=bool(rewrite_prompt),
        api_name="/encode",
    )

    with safe_open(path, framework="pt") as handle:
        metadata = handle.metadata()

        return (
            handle.get_tensor("prompt_embeds"),
            handle.get_tensor("text_token_tags"),
            metadata,
            plan,
        )


# ----------------------------------------------------------------------
# GPU DURATION ESTIMATION
# ----------------------------------------------------------------------

_DUR_B = 1.1745e-4
_DUR_C = 3.8396e-9

_DECODE_BASE = 15
_DECODE_PER_DEFAULT_CANVAS = 15
_DEFAULT_CANVAS_PIXELS = 960 * 544 * 124

_PLACEMENT_ALLOWANCE = 12
_PAD = 10


def get_duration(
    prompt_embeds,
    text_token_tags,
    image,
    last_image,
    height,
    width,
    num_frames,
    steps,
    seed,
    lora="larry",
    *a,
    **k,
):
    height = int(height)
    width = int(width)
    num_frames = int(num_frames)
    steps = int(steps)

    latent_frames = (
        (num_frames - LATENTS_PER_CHUNK)
        // FRAMES_PER_CHUNK
        * LATENTS_PER_CHUNK
        + 2
    )

    patches = (height // 32) * (width // 32)

    rows = (
        latent_frames * patches
        + (
            int(image is not None)
            + int(last_image is not None)
        )
        * patches
    )

    denoise = steps * (
        _DUR_B * rows
        + _DUR_C * rows**2
    )

    decode = (
        _DECODE_BASE
        + _DECODE_PER_DEFAULT_CANVAS
        * (height * width * num_frames)
        / _DEFAULT_CANVAS_PIXELS
    )

    return max(
        60,
        int(denoise + decode)
        + _PLACEMENT_ALLOWANCE
        + _PAD,
    )


# ----------------------------------------------------------------------
# GPU GENERATION
# ----------------------------------------------------------------------

@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="larry",
):
    """Run the denoise loop and decoders on GPU."""

    import torch
    import h3_lora

    active_lora = h3_lora.set_active(
        PIPE.transformer,
        lora,
    )

    if PLACEMENT == "lazy":
        PIPE.to("cuda")

    elif PLACEMENT == "pack":
        PIPE.vae.to("cuda")
        PIPE.audio_vae.to("cuda")

    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],
        state.get("audio")[0].cpu(),
        state.get("sampling_rate"),
        active_lora,
    )


# ----------------------------------------------------------------------
# KEYFRAME FITTING
# ----------------------------------------------------------------------

def _fit_keyframe(
    image_path,
    current_canvas,
):
    """Cover-crop an uploaded keyframe to a supported aspect ratio."""

    from PIL import Image as _Image

    img = _Image.open(image_path)

    aspect = img.width / img.height

    fastest = {}

    for label, (h, w) in CANVASES.items():
        r = w / h

        if (
            r not in fastest
            or w * h
            < fastest[r][1][0] * fastest[r][1][1]
        ):
            fastest[r] = (
                label,
                (h, w),
            )

    ratio = min(
        fastest,
        key=lambda r: abs(r - aspect),
    )

    label, (h, w) = fastest[ratio]

    cur_h, cur_w = CANVASES[current_canvas]

    if abs(
        cur_w / cur_h - aspect
    ) <= abs(ratio - aspect):
        label = current_canvas
        h, w = cur_h, cur_w

    target = w / h

    if abs(
        img.width / img.height - target
    ) > 1e-3:

        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, label


# ----------------------------------------------------------------------
# LORA
# ----------------------------------------------------------------------

def _resolve_lora(
    lora,
    use_lora,
) -> str:
    """Resolve the requested LoRA."""

    if isinstance(lora, str) and lora in (
        "larry",
        "lightx",
        "off",
    ):
        return lora

    return (
        "larry"
        if use_lora
        else "off"
    )


# ----------------------------------------------------------------------
# GENERATION FUNCTION
# ----------------------------------------------------------------------

def generate(
    prompt,
    image_path=None,
    last_image_path=None,
    canvas=DEFAULT_CANVAS,
    duration=5,
    steps=6,
    seed=42,
    upsample=False,
    use_lora=True,
    lora="",
    ip_token=None,
):
    """Generate one video and archive a private copy."""

    if LOAD_ERROR:
        raise Exception(LOAD_ERROR)

    if PIPE is None:
        raise Exception(
            "The denoiser is still loading."
        )

    if not prompt or not prompt.strip():
        raise Exception(
            "MiniMax-H3 always takes a prompt, "
            "keyframes or not."
        )

    from PIL import Image, ImageOps
    from diffusers.utils import encode_video

    # Resolve LoRA.
    lora = _resolve_lora(
        lora,
        use_lora,
    )

    # --------------------------------------------------------------
    # Resolve uploaded keyframes.
    # --------------------------------------------------------------

    first = (
        image_path["path"]
        if isinstance(image_path, dict)
        else image_path
    )

    last = (
        last_image_path["path"]
        if isinstance(last_image_path, dict)
        else last_image_path
    )

    if first:
        first, canvas = _fit_keyframe(
            first,
            canvas,
        )

    if last:
        last, canvas = _fit_keyframe(
            last,
            canvas,
        )

    # --------------------------------------------------------------
    # Calculate frame count.
    # --------------------------------------------------------------

    num_frames = snap_frames(
        duration
    )

    # --------------------------------------------------------------
    # Conditioner.
    # --------------------------------------------------------------

    conditioned = time.time()

    (
        prompt_embeds,
        text_token_tags,
        metadata,
        plan,
    ) = encode_remote(
        prompt,
        first,
        last,
        canvas,
        num_frames,
        rewrite_prompt=upsample,
        ip_token=ip_token,
    )

    condition_seconds = (
        time.time() - conditioned
    )

    height, width, num_frames = (
        int(metadata[key])
        for key in (
            "height",
            "width",
            "num_frames",
        )
    )

    refined = (
        plan.get("refined_prompt")
        or ""
    )

    # --------------------------------------------------------------
    # Convert keyframe to RGB.
    # --------------------------------------------------------------

    def keyframe(path):
        return (
            ImageOps.exif_transpose(
                Image.open(path)
            ).convert("RGB")
            if path
            else None
        )

    # --------------------------------------------------------------
    # GPU generation.
    # --------------------------------------------------------------

    started = time.time()

    (
        frames,
        audio,
        sampling_rate,
        active_lora,
    ) = _generate(
        prompt_embeds,
        text_token_tags,
        keyframe(first),
        keyframe(last),
        height,
        width,
        num_frames,
        steps,
        seed,
        lora,
    )

    generate_seconds = (
        time.time() - started
    )

    # --------------------------------------------------------------
    # Make sure both directories exist.
    # --------------------------------------------------------------

    os.makedirs(
        OUTPUT_DIR,
        exist_ok=True,
    )

    os.makedirs(
        PRIVATE_OUTPUT_DIR,
        exist_ok=True,
    )

    # --------------------------------------------------------------
    # Unique generation ID.
    # --------------------------------------------------------------

    generation_id = (
        f"h3-{int(time.time() * 1000)}"
    )

    # --------------------------------------------------------------
    # PUBLIC / TEMPORARY OUTPUT
    #
    # This is the copy returned to the user.
    # --------------------------------------------------------------

    path = os.path.join(
        OUTPUT_DIR,
        f"{generation_id}.mp4",
    )

    encode_video(
        frames,
        fps=FPS,
        output_path=path,
        audio=audio,
        audio_sample_rate=sampling_rate,
    )

    # --------------------------------------------------------------
    # PRIVATE PERSISTENT BACKUP
    #
    # This copy goes into the mounted private bucket.
    #
    # IMPORTANT:
    # We do NOT return this path to Gradio.
    # --------------------------------------------------------------

    private_path = os.path.join(
        PRIVATE_OUTPUT_DIR,
        f"{generation_id}.mp4",
    )

    shutil.copy2(
        path,
        private_path,
    )

    # --------------------------------------------------------------
    # PRIVATE METADATA BACKUP
    # --------------------------------------------------------------

    metadata_path = os.path.join(
        PRIVATE_OUTPUT_DIR,
        f"{generation_id}.json",
    )

    metadata_record = {
        "timestamp": time.time(),
        "generation_id": generation_id,
        "prompt": prompt,
        "refined_prompt": refined,
        "width": width,
        "height": height,
        "frames": num_frames,
        "duration_seconds": (
            num_frames / FPS
        ),
        "steps": int(steps),
        "seed": int(seed),
        "lora": active_lora,
        "canvas": canvas,
        "model_repo": MODEL_REPO,
        "conditioner_space": CONDITIONER_SPACE,
    }

    with open(
        metadata_path,
        "w",
        encoding="utf-8",
    ) as metadata_file:
        json.dump(
            metadata_record,
            metadata_file,
            indent=2,
            ensure_ascii=False,
        )

    # --------------------------------------------------------------
    # Report shown to the user.
    # --------------------------------------------------------------

    report = (
        f"{width}x{height} · "
        f"{num_frames} frames "
        f"({num_frames / FPS:.3f} s) · "
        f"{int(steps)} steps · "
        f"conditioner "
        f"{condition_seconds:.0f}s "
        f"({plan['num_text_tokens']} tokens"
        f"{', upsampled' if refined else ''}) · "
        f"denoise + decode "
        f"{generate_seconds:.0f}s "
        f"({generate_seconds / int(steps):.1f} s/step) · "
        f"turbo LoRA {active_lora} · "
        f"seed {int(seed)}"
    )

    print(
        f"[gen] {report}",
        flush=True,
    )

    print(
        f"[archive] private video: {private_path}",
        flush=True,
    )

    print(
        f"[archive] private metadata: {metadata_path}",
        flush=True,
    )

    # Return ONLY the temporary/public copy.
    return (
        FileData(path=path),
        report,
        refined,
    )


# ======================================================================
# SERVER MODE
# ======================================================================

app = Server(
    title="MiniMax-H3 Studio"
)


# ----------------------------------------------------------------------
# GENERATE API
# ----------------------------------------------------------------------

@app.api(name="generate")
def _generate_api(
    prompt: str,
    image_path: FileData | None = None,
    last_image_path: FileData | None = None,
    canvas: str = DEFAULT_CANVAS,
    duration: float = 5,
    steps: int = 6,
    seed: float = 42,
    upsample: bool = False,
    use_lora: bool = True,
    lora: str = "",
    request: Request = None,
) -> tuple[FileData, str, str]:

    """Generate a video with synchronized soundtrack."""

    # The request's x-ip-token bills the conditioner
    # to the caller when available.
    ip_token = (
        request.headers.get("x-ip-token")
        if request is not None
        else None
    )

    return generate(
        prompt,
        image_path,
        last_image_path,
        canvas,
        duration,
        steps,
        seed,
        upsample,
        use_lora,
        lora,
        ip_token=ip_token,
    )


# ----------------------------------------------------------------------
# STATUS
# ----------------------------------------------------------------------

@app.get("/status")
def studio_status():
    """Return model readiness."""

    return {
        "ready": (
            PIPE is not None
            and LOAD_ERROR is None
        ),
        "status": status(),
    }


# ----------------------------------------------------------------------
# STUDIO CONFIG
# ----------------------------------------------------------------------

@app.get("/studio-config")
def studio_config():
    """Return canvas and LoRA configuration."""

    import h3_lora

    state = (
        getattr(
            PIPE.transformer,
            "_lora_state",
            None,
        )
        if PIPE is not None
        else None
    )

    sets = (
        state["sets"]
        if state
        else {}
    )

    return {
        "canvases": list(CANVASES),
        "default_canvas": DEFAULT_CANVAS,
        "min_duration": MIN_UI_DURATION,
        "max_duration": MAX_UI_DURATION,

        "loras": {
            **{
                name: {
                    "label": spec["label"],
                    "steps": {
                        "larry": 6,
                        "lightx": 4,
                    }.get(
                        name,
                        6,
                    ),
                }
                for name, spec in sets.items()
            },

            "off": {
                "label": "off (base model)",
                "steps": 28,
            },
        },

        "default_lora": (
            state["active"]
            if state
            else "off"
        ),
    }


# ----------------------------------------------------------------------
# HOMEPAGE
# ----------------------------------------------------------------------

@app.get(
    "/",
    response_class=HTMLResponse,
)
def homepage():

    with open(
        os.path.join(
            os.path.dirname(
                os.path.abspath(__file__)
            ),
            "index.html",
        ),
        encoding="utf-8",
    ) as f:
        return f.read()


# ----------------------------------------------------------------------
# LOAD MODELS
# ----------------------------------------------------------------------

load_models()


# ----------------------------------------------------------------------
# START SERVER
# ----------------------------------------------------------------------

if __name__ == "__main__":

    # IMPORTANT:
    # Only OUTPUT_DIR is included here.
    #
    # DO NOT add PRIVATE_OUTPUT_DIR.
    #
    # This prevents the private archive from being intentionally exposed
    # through Gradio's /gradio_api/file= route.
    app.launch(
        show_error=True,
        allowed_paths=[
            OUTPUT_DIR
        ],
    )