from __future__ import annotations

import os
import tempfile
import time
from functools import lru_cache

import spaces

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

import gradio as gr
import torch

from h3_split_blocks import MiniMaxH3GeneratorBlocks


MODEL_REPO = "MiniMaxAI/MiniMax-H3"
LORA_REPO = "rzgar/minimax-h3_fl2v_8Step_motion_enhancer"
LORA_FILE = "minimax-h3_fl2v_8Step_motion_enhancer.safetensors"
CONDITIONER_SPACE = "multimodalart/qwen3vl-conditioner"
FPS = 24
STEPS = 8
OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "minimax-h3-motion-enhancer")
os.makedirs(OUTPUT_DIR, exist_ok=True)

CANVASES = {
    "960x544 · 16:9 fast": (544, 960),
    "1024x576 · 16:9 fast": (576, 1024),
    "544x960 · 9:16 fast": (960, 544),
    "544x544 · 1:1 fast": (544, 544),
    "768x576 · 4:3 fast": (576, 768),
    "576x768 · 3:4 fast": (768, 576),
    "1152x512 · 21:9 fast": (512, 1152),
}
DEFAULT_CANVAS = "960x544 · 16:9 fast"


def _reorder_interleaved_qkv(weight: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor:
    grouped = weight.reshape(num_heads, 3 * head_dim, *weight.shape[1:])
    query, key, value = grouped.split(head_dim, dim=1)
    return torch.cat(
        [part.reshape(num_heads * head_dim, *weight.shape[1:]) for part in (query, key, value)], dim=0
    )


def _lora_target_name(source_name: str) -> str:
    target = source_name.removeprefix("diffusion_model.")
    if target.startswith("token_refiner.blocks."):
        target = target.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
    elif target.startswith("blocks."):
        target = target.replace("blocks.", "transformer_blocks.", 1)
    return target.replace("final_layer.adaln_proj.linear", "norm_out.linear")


def _lora_targets(name: str, b_weight: torch.Tensor, num_heads: int, head_dim: int):
    target = _lora_target_name(name)
    if target.endswith(".attn.qkv_proj"):
        prefix = target.removesuffix("qkv_proj")
        rows = _reorder_interleaved_qkv(b_weight, num_heads, head_dim)
        for kind, part in zip(("q", "k", "v"), rows.split(num_heads * head_dim, dim=0)):
            yield f"{prefix}to_{kind}.weight", part.contiguous()
    elif target.endswith(".mlp.fc1"):
        gate, value = b_weight.chunk(2, dim=0)
        yield target.replace(".mlp.fc1", ".ff.net.0.proj") + ".weight", torch.cat([value, gate]).contiguous()
    elif target.endswith(".mlp.fc2"):
        yield target.replace(".mlp.fc2", ".ff.net.2") + ".weight", b_weight
    elif target.endswith(".attn.out_proj"):
        yield target.replace(".attn.out_proj", ".attn.to_out.0") + ".weight", b_weight
    else:
        yield target + ".weight", b_weight


def _apply_motion_lora(transformer) -> str:
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    path = hf_hub_download(LORA_REPO, LORA_FILE)
    lora = load_file(path, device="cpu")
    suffix_a = ".lora_A.weight"
    suffix_b = ".lora_B.weight"
    bases = sorted({key[: -len(suffix_a)] for key in lora if key.endswith(suffix_a)})
    if not bases:
        raise ValueError("The motion enhancer LoRA did not contain any lora_A/lora_B pairs")

    params = dict(transformer.named_parameters())
    num_heads = transformer.config.num_attention_heads
    head_dim = transformer.config.attention_head_dim
    folded = 0
    missed = []
    with torch.no_grad():
        for name in bases:
            a = lora[f"{name}{suffix_a}"]
            b = lora.get(f"{name}{suffix_b}")
            if b is None:
                raise ValueError(f"Missing LoRA B matrix for {name}")
            for key, b_part in _lora_targets(name, b, num_heads, head_dim):
                param = params.get(key)
                if param is None:
                    missed.append(key)
                    continue
                delta = b_part.float() @ a.float()
                if delta.shape != param.shape:
                    raise ValueError(f"LoRA shape mismatch for {key}: {delta.shape} vs {param.shape}")
                param.data = (param.data.float() + delta.to(param.device)).to(param.dtype)
                folded += 1

    if missed:
        raise ValueError(f"{len(missed)} LoRA targets did not match the H3 transformer, e.g. {missed[:3]}")
    return f"Motion enhancer loaded: {len(bases)} adapters, {folded} matrices"


def _load_pipeline():
    from diffusers import ComponentsManager

    manager = ComponentsManager()
    blocks = MiniMaxH3GeneratorBlocks()
    pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
    pipe.load_components(dtype=torch.bfloat16)
    pipe.vae.set_attention_backend("native")
    pipe.audio_vae.set_attention_backend("native")
    pipe.transformer.set_attention_backend(os.environ.get("H3_ATTENTION", "_native_cudnn"))
    print(_apply_motion_lora(pipe.transformer), flush=True)
    # Module-scope placement lets ZeroGPU pack the weights once instead of per request.
    pipe.transformer.to("cuda")
    return pipe


PIPE = _load_pipeline()


@lru_cache(maxsize=1)
def _conditioner_client():
    from gradio_client import Client

    return Client(CONDITIONER_SPACE)


def _snap_frames(seconds: float) -> int:
    frames = max(120, round(float(seconds) * FPS))
    while frames % 17 != 5:
        frames += 1
    return frames


def _cover_crop(image_path: str, canvas: str) -> str:
    from PIL import Image, ImageOps

    height, width = CANVASES[canvas]
    with Image.open(image_path) as image:
        image = ImageOps.exif_transpose(image).convert("RGB")
        image = ImageOps.fit(image, (width, height), method=Image.Resampling.LANCZOS)
        output = os.path.join(OUTPUT_DIR, f"frame-{time.time_ns()}.png")
        image.save(output)
    return output


def _condition(prompt: str, image_path: str, canvas: str, num_frames: int, rewrite_prompt: bool = False):
    """Encode the prompt through the same Qwen3-VL conditioner path used by FlashGen.

    When rewrite_prompt=True, the conditioner first rewrites the user's prompt into
    MiniMax-H3's structured shot description, then encodes that rewritten prompt.
    This is the "Upsample/Enhance prompt" behavior from the FlashGen Space.
    """
    from gradio_client import handle_file
    from safetensors import safe_open

    def call():
        return _conditioner_client().predict(
            prompt=prompt,
            image_path=handle_file(image_path),
            last_image_path=None,
            canvas=canvas,
            num_frames=num_frames,
            rewrite_prompt=bool(rewrite_prompt),
            api_name="/encode",
        )

    try:
        path, plan = call()
    except Exception as first:
        # The conditioner can occasionally leave a stale Gradio client behind.
        # Retry once with a fresh client, matching the FlashGen implementation.
        print(f"[conditioner] retrying with a fresh client after: {first}", flush=True)
        _conditioner_client.cache_clear()
        path, plan = call()

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


@spaces.GPU(duration=120, size="xlarge")
def _generate_gpu(embeds, tags, image, height: int, width: int, num_frames: int, seed: int):
    """Run MiniMax-H3 denoising and return CPU video frames plus soundtrack."""
    PIPE.vae.to("cuda")
    PIPE.audio_vae.to("cuda")
    pipe_state = PIPE(
        prompt_embeds=embeds.to("cuda"),
        text_token_tags=tags,
        image=image,
        height=int(height),
        width=int(width),
        num_frames=int(num_frames),
        num_inference_steps=STEPS + 1,
        generator=torch.Generator("cpu").manual_seed(int(seed)),
    )
    frames = pipe_state.get("videos")[0]
    audio = pipe_state.get("audio")[0]
    if torch.is_tensor(frames):
        frames = frames.cpu()
    if torch.is_tensor(audio):
        audio = audio.cpu()
    return frames, audio, int(pipe_state.get("sampling_rate"))


def generate(
    prompt: str,
    image_path: str | None,
    canvas: str,
    duration: float,
    seed: int,
    enhance_prompt: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Animate an uploaded image, optionally using H3's Qwen3-VL prompt enhancement."""
    if not image_path:
        raise gr.Error("Upload a first frame before generating.")
    if not prompt or not prompt.strip():
        raise gr.Error("Describe the motion you want to see.")

    height, width = CANVASES[canvas]
    num_frames = _snap_frames(duration)
    cropped_path = _cover_crop(image_path, canvas)
    from PIL import Image
    from diffusers.utils import encode_video

    progress(
        0.05,
        desc="Enhancing prompt, then encoding..." if enhance_prompt else "Encoding the first frame...",
    )
    embeds, tags, metadata, plan = _condition(
        prompt.strip(), cropped_path, canvas, num_frames, rewrite_prompt=bool(enhance_prompt)
    )
    height = int(metadata.get("height", height))
    width = int(metadata.get("width", width))
    num_frames = int(metadata.get("num_frames", num_frames))

    enhanced = (plan or {}).get("refined_prompt") or ""
    progress(0.15, desc=f"Generating {num_frames / FPS:.1f}s at {width}x{height} in {STEPS} steps")
    with Image.open(cropped_path) as image:
        frames, audio, sampling_rate = _generate_gpu(
            embeds, tags, image.convert("RGB"), height, width, num_frames, int(seed)
        )

    output = os.path.join(OUTPUT_DIR, f"motion-enhanced-{time.time_ns()}.mp4")
    encode_video(frames, fps=FPS, output_path=output, audio=audio, audio_sample_rate=sampling_rate)

    status = (
        f"{width}x{height} · {num_frames / FPS:.2f}s · {STEPS} denoising steps · "
        f"seed {int(seed)}"
    )
    return output, status, enhanced, gr.update(visible=bool(enhanced))


DESCRIPTION = """
# MiniMax-H3 Motion Enhancer

Upload a still image and describe the movement. The Space uses the **8-step Motion Enhancer LoRA**
with MiniMax-H3 and synchronized audio.

**Enhance / Upsample prompt** uses the same Qwen3-VL prompt-rewriting path as the MiniMax-H3 FlashGen
Space: your short instruction is first expanded into a structured shot/action/sound description and
then those rewritten embeddings are sent to H3.

Use only images you own or have permission to animate.
"""


with gr.Blocks(title="MiniMax-H3 Motion Enhancer") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            image = gr.Image(label="First frame", type="filepath", sources=["upload", "webcam"])
            prompt = gr.Textbox(
                label="Motion prompt",
                value="The person naturally performs the described action. Keep the same identity and appearance. Natural body motion, realistic facial expression, synchronized sound effects and speech if specified.",
                lines=5,
                placeholder="Describe the action, camera movement and sound...",
            )
            enhance = gr.Checkbox(
                label="Enhance / Upsample prompt",
                value=False,
                info="Rewrite the prompt with H3's Qwen3-VL into a detailed structured shot description before generation.",
            )
            with gr.Row():
                canvas = gr.Dropdown(list(CANVASES), value=DEFAULT_CANVAS, label="Output canvas")
                duration = gr.Slider(5, 15, value=5, step=1, label="Duration (seconds)")
            seed = gr.Number(value=42, precision=0, label="Seed")
            generate_button = gr.Button("Generate motion", variant="primary", size="lg")

        with gr.Column(scale=1):
            video = gr.Video(label="Generated video + soundtrack", autoplay=True)
            status = gr.Markdown()
            with gr.Accordion("Upsampled prompt", open=False, visible=False) as enhanced_panel:
                enhanced_box = gr.Textbox(show_label=False, lines=10, interactive=False)

    outputs = [video, status, enhanced_box, enhanced_panel]

    generate_button.click(
        fn=generate,
        inputs=[prompt, image, canvas, duration, seed, enhance],
        outputs=outputs,
        api_name="generate",
    )

    prompt.submit(
        fn=generate,
        inputs=[prompt, image, canvas, duration, seed, enhance],
        outputs=outputs,
        api_name="generate_submit",
    )


demo.queue(max_size=8).launch(mcp_server=True)