"""MiniMax-H3 Ref2VA Generator — denoise half (omni references).

Run `../scripts/sync_diffusers.sh generator_ref2va` before deploy.

API `/generate` accepts either:
- Orchestrator style: `ref_images` / `ref_audios` / `ref_videos` lists
- multimodalart style: `image_1`…`image_9`, `audio_path`, `video_path`
"""

from __future__ import annotations

import os
import tempfile
import time
import traceback

import spaces  # noqa: F401
import gradio as gr

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
CONDITIONER = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn")
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
# ZeroGPU preflight; actual Ref2VA ~4–8 min. Keep reservation tight.
# ZeroGPU preflight ask ≈ 2× this value. Keep ≥300 so ask≈600 for Ref2VA.
GPU_DURATION = int(os.environ.get("H3_GPU_DURATION", os.environ.get("GPU_DURATION", "300")))
FPS = 24
MIN_REF_VIDEO, MAX_REF_VIDEO = 2.0, 15.0

CANVASES = {
    "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),
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    "1152x512 · 21:9 fast": (512, 1152),
    "1536x672 · 21:9 full": (672, 1536),
}
DEFAULT_CANVAS = "960x544 · 16:9 fast"

PIPE = None
LOAD_ERROR: str | None = None
CLIENT = None


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


def _as_paths(items) -> list[str]:
    out = []
    for item in items or []:
        if not item:
            continue
        out.append(item if isinstance(item, str) else getattr(item, "name", None) or str(item))
    return out


def collect(images, audios, videos) -> list[tuple[str, str]]:
    """Order: images → audios → videos (matches multimodalart / model presentation)."""
    ordered = [("image", p) for p in _as_paths(images)]
    ordered += [("audio", p) for p in _as_paths(audios)]
    ordered += [("video", p) for p in _as_paths(videos)]
    return ordered


def probe(path: str):
    import av

    with av.open(path) as container:
        def seconds(stream, cont):
            if stream.duration is not None and stream.time_base is not None:
                return float(stream.duration * stream.time_base)
            if cont.duration is not None:
                return cont.duration / 1_000_000
            return None

        video = seconds(container.streams.video[0], container) if container.streams.video else None
        audio = seconds(container.streams.audio[0], container) if container.streams.audio else None
    return video, audio


def check(prompt: str, references: list[tuple[str, str]]) -> None:
    if not prompt or not prompt.strip():
        raise gr.Error("prompt required")
    if not references:
        raise gr.Error("Add at least one reference image or video")
    if {k for k, _ in references} == {"audio"}:
        raise gr.Error("Audio cannot be the sole reference")
    for kind, path in references:
        if kind != "video":
            continue
        video_s, _ = probe(path)
        if video_s is None:
            raise gr.Error(f"No video stream in {path}")
        if not MIN_REF_VIDEO <= video_s <= MAX_REF_VIDEO:
            raise gr.Error(f"Reference video {video_s:.1f}s out of {MIN_REF_VIDEO}-{MAX_REF_VIDEO}s")


def load_models() -> str | None:
    global PIPE, LOAD_ERROR
    if PIPE is not None or LOAD_ERROR is not None:
        return LOAD_ERROR
    token = os.environ.get("HF_TOKEN")
    if not token:
        LOAD_ERROR = "**`HF_TOKEN` missing**"
        return LOAD_ERROR
    try:
        import torch
        from diffusers import ComponentsManager
        from h3_split_blocks import MiniMaxH3Ref2VAGeneratorBlocks
        import h3_aoti

        manager = ComponentsManager()
        blocks = MiniMaxH3Ref2VAGeneratorBlocks()
        print(f"[ref2va] 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, token=token)
        pipe.transformer_ref.set_attention_backend(ATTENTION)
        h3_aoti.maybe_load(pipe.transformer_ref)
        PIPE = pipe
        print("[ref2va] ready", flush=True)
    except Exception as err:
        traceback.print_exc()
        LOAD_ERROR = (
            f"**Ref2VA load failed:** `{type(err).__name__}: {err}` — "
            "run `scripts/sync_diffusers.sh generator_ref2va`."
        )
    return LOAD_ERROR


def conditioner():
    global CLIENT
    if CLIENT is None:
        from gradio_client import Client

        CLIENT = Client(CONDITIONER)
    return CLIENT


def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
    from gradio_client import handle_file
    from safetensors import safe_open

    # multimodalart/qwen3vl-conditioner requires rewrite_prompt (6th arg).
    # Older conditioners (5-arg) reject unknown kwargs — fall back.
    kwargs = dict(
        prompt=prompt,
        media=[handle_file(p) for _, p in references],
        kinds=",".join(k for k, _ in references),
        canvas=canvas,
        num_frames=num_frames,
        api_name="/encode_ref2va",
    )
    try:
        # Current multimodalart/qwen3vl-conditioner: 6 params including rewrite_prompt.
        path, plan = conditioner().predict(**kwargs, rewrite_prompt=bool(rewrite_prompt))
    except Exception:
        # Older 5-arg conditioners (e.g. owenisas/minimax-h3-conditioner).
        path, plan = conditioner().predict(**kwargs)
    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
def _generate(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed):
    import torch
    from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3Reference

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

    state = PIPE(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        references=[MiniMaxH3Reference(**{kind: path}) for kind, path in references],
        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")


def _run(prompt, references, canvas, duration, steps, seed, progress):
    from diffusers.utils import encode_video

    try:
        check(prompt, references)
        requested = snap_frames(duration)
        progress(0.0, desc=f"Conditioning via {CONDITIONER}")
        embeds, tags, meta, plan = encode_remote(prompt, references, canvas, requested)
        height, width, num_frames = (int(meta[k]) for k in ("height", "width", "num_frames"))
        progress(0.15, desc="Denoising Ref2VA")
        t0 = time.time()
        frames, audio, sr = _generate(embeds, tags, references, height, width, num_frames, steps, seed)
        elapsed = time.time() - t0
        out_dir = os.path.join(tempfile.gettempdir(), "h3-ref2va")
        os.makedirs(out_dir, exist_ok=True)
        path = os.path.join(out_dir, f"h3-ref2va-{int(time.time() * 1000)}.mp4")
        encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sr)
        report = (
            f"refs={[k for k, _ in references]} · `{width}x{height}`, {num_frames} frames, {steps} steps · "
            f"denoise {elapsed:.0f}s ({elapsed / max(int(steps), 1):.1f}s/step) · seed {seed}"
        )
        return path, report
    except gr.Error:
        raise
    except Exception as err:
        traceback.print_exc()
        raise gr.Error(f"{type(err).__name__}: {err}")


def generate(
    prompt,
    image_1=None,
    audio_path=None,
    video_path=None,
    canvas=DEFAULT_CANVAS,
    image_2=None,
    image_3=None,
    image_4=None,
    image_5=None,
    image_6=None,
    image_7=None,
    image_8=None,
    image_9=None,
    duration=5,
    steps=28,
    seed=42,
    ref_images=None,
    ref_audios=None,
    ref_videos=None,
    progress=gr.Progress(track_tqdm=True),
):
    """Compatible with multimodalart positional slots + our list-based orchestrator fields."""
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("Generator still loading")

    if ref_images or ref_audios or ref_videos:
        references = collect(ref_images, ref_audios, ref_videos)
    else:
        images = [image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9]
        references = collect(images, [audio_path] if audio_path else [], [video_path] if video_path else [])
    return _run(prompt, references, canvas, duration, steps, seed, progress)


load_models()

with gr.Blocks(title="MiniMax-H3 Ref2VA") as demo:
    gr.Markdown(
        f"# MiniMax-H3 Ref2VA\nOmni references → video + stereo audio. Conditioner: `{CONDITIONER}`."
    )
    prompt = gr.Textbox(label="Prompt", lines=3)
    ref_images = gr.File(label="Reference images (≤9)", file_count="multiple", type="filepath", file_types=["image"])
    ref_audios = gr.File(label="Reference audio (≤3)", file_count="multiple", type="filepath", file_types=["audio"])
    ref_videos = gr.File(label="Reference videos (≤3)", file_count="multiple", type="filepath", file_types=["video"])
    canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
    duration = gr.Slider(2, 14, value=5, step=1, label="Duration (s)")
    steps = gr.Slider(10, 40, value=28, step=1, label="Steps")
    seed = gr.Number(value=42, precision=0, label="Seed")
    btn = gr.Button("Generate", variant="primary")
    video = gr.Video()
    report = gr.Markdown()

    def ui_generate(prompt, ref_images, ref_audios, ref_videos, canvas, duration, steps, seed, progress=gr.Progress()):
        return generate(
            prompt=prompt,
            canvas=canvas,
            duration=duration,
            steps=steps,
            seed=seed,
            ref_images=ref_images,
            ref_audios=ref_audios,
            ref_videos=ref_videos,
            progress=progress,
        )

    btn.click(
        ui_generate,
        [prompt, ref_images, ref_audios, ref_videos, canvas, duration, steps, seed],
        [video, report],
        api_name="generate",
    )


if __name__ == "__main__":
    demo.launch(show_error=True)