"""A MiniMax-H3 modular workflow that accepts keyframes **and** references in the same run.

Why this exists. MiniMax-H3 denoises one packed sequence, and that sequence can hold keyframe conditioning rows and
reference conditioning rows at the same time. `diffusers`' shipped blocks cannot express it: their conditional steps
dispatch either/or (`select_block` checks `references` first), so a request carrying both is accepted and the
keyframes are **silently dropped** — no error, no warning, just a reference-only generation.

`MiniMaxH3CombinedBlocks` is `MiniMaxH3Blocks` with three of its conditional steps replaced by ones that know a
fourth shape. Nothing else changes: `t2va`, `fl2va` and `ref2va` requests take exactly the same path they always did,
and the denoising loop is untouched, because the conditioning rows are simply the leading rows of the sequence and it
only ever steps what comes after them.

    from diffusers import ModularPipeline

    pipe = ModularPipeline.from_pretrained(REPO, trust_remote_code=True, workflow="combined")
    pipe.load_components(dtype=torch.bfloat16, trust_remote_code=True)
    out = pipe(prompt=..., references=[...], image=first, last_image=last, num_frames=124, ...)

The layout itself lives in `combined_layout.py`, which is pinned by reproducing both shipped builders bit for bit in
their degenerate cases.
"""

from __future__ import annotations

import numpy as np
import torch
from diffusers.modular_pipelines import ConditionalPipelineBlocks, SequentialPipelineBlocks
from diffusers.modular_pipelines.minimax_h3.before_denoise import (
    MiniMaxH3PrepareConditionLatentsStep,
    MiniMaxH3PrepareLatentsStep,
    MiniMaxH3Ref2VAPrepareLatentsStep,
    MiniMaxH3Ref2VAPrepareLayoutStep,
    MiniMaxH3SetTimestepsStep,
)
from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
from diffusers.modular_pipelines.minimax_h3.decoders import MiniMaxH3AfterDenoiseStep
from diffusers.modular_pipelines.minimax_h3.denoise import MiniMaxH3Ref2VADenoiseStep
from diffusers.modular_pipelines.minimax_h3.encoders import (
    MiniMaxH3Ref2VAReferenceEncoderStep,
    encode_vae_condition,
)
from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
    MiniMaxH3AutoDenoiseStep,
    MiniMaxH3AutoTextEncoderStep,
    MiniMaxH3AutoVaeEncoderStep,
    MiniMaxH3AutoBeforeEncodeStep,
    MiniMaxH3Blocks,
    MiniMaxH3DecodeStep,
)
from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam

from .combined_layout import build_combined_packed_sequence


def _anchors_of(image, last_image) -> tuple[str, ...]:
    """Which end of the clip each keyframe is anchored to, in packed order."""
    return tuple(name for name, value in (("first", image), ("last", last_image)) if value is not None)


class MiniMaxH3KeyframesOnCanvasStep(ModularPipelineBlocks):
    """Put the keyframes on the canvas the reference setup already resolved.

    The `fl2va` resize step *derives* the canvas from the keyframe's aspect ratio. Here the references have already
    settled it, so the keyframes are stretched onto it — which is what `fl2va` does to a keyframe whose aspect does
    not match the target anyway.
    """

    model_name = "minimax-h3"

    @property
    def description(self) -> str:
        return "Stretches the keyframes of a combined request onto the canvas the reference setup resolved."

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(name="image", description="Keyframe the video starts from."),
            InputParam(name="last_image", description="Keyframe the video ends on."),
            InputParam(name="height", type_hint=int, required=True, description="Resolved height in pixels."),
            InputParam(name="width", type_hint=int, required=True, description="Resolved width in pixels."),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("keyframes", type_hint=list, description="The keyframes on the target canvas, packed order."),
            OutputParam("keyframe_anchors", type_hint=tuple, description="Which end each keyframe is anchored to."),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        size = (block_state.width, block_state.height)
        block_state.keyframe_anchors = _anchors_of(block_state.image, block_state.last_image)
        block_state.keyframes = [
            frame.convert("RGB").resize(size)
            for frame in (block_state.image, block_state.last_image)
            if frame is not None
        ]
        self.set_block_state(state, block_state)
        return components, state


class MiniMaxH3CombinedKeyframeEncoderStep(ModularPipelineBlocks):
    """Encode the keyframes and put them **in front of** the reference latents.

    Order is the contract: the combined layout reserves `[text | keyframe cond | reference blocks | targets]`, and the
    stock prepare-latents step concatenates this list in order — then asserts the rows it produced equal the rows the
    layout reserved, so a mismatch raises instead of degrading quietly.
    """

    model_name = "minimax-h3"

    @property
    def description(self) -> str:
        return "Encodes a combined request's keyframes and prepends them to the reference conditioning latents."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("vae")]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(name="keyframes", type_hint=list, required=True, description="Keyframes on the canvas."),
            InputParam(name="condition_latents", type_hint=list, required=True,
                       description="Reference conditioning latents, packed order."),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("condition_latents", type_hint=list,
                        description="Keyframe latents first, then the reference latents."),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = components._execution_device
        keyframe_latents = [
            encode_vae_condition(
                components.vae,
                torch.from_numpy(np.array(image)).to(device).permute(2, 0, 1)[None, :, None],
                components.pixel_mean,
                components.pixel_std,
                components.keyframe_encode_seed,
            )
            for image in block_state.keyframes
        ]
        block_state.condition_latents = keyframe_latents + list(block_state.condition_latents)
        self.set_block_state(state, block_state)
        return components, state


class MiniMaxH3CombinedPrepareLayoutStep(MiniMaxH3Ref2VAPrepareLayoutStep):
    """The `ref2va` layout step, with keyframe rows packed ahead of the reference blocks."""

    @property
    def description(self) -> str:
        return (
            "Resolves the latent shapes of a combined request and builds its packed layout — "
            "`[text | keyframe conditions | reference blocks | target audio | target video]`. The references push the "
            "target timeline out, so the keyframe anchors ride on the timeline their spans leave behind."
        )

    @property
    def inputs(self) -> list[InputParam]:
        return super().inputs + [
            InputParam(name="keyframe_anchors", type_hint=tuple, default=(),
                       description="Which end of the video each keyframe is anchored to, in packed order."),
        ]

    # Overrides the parent's `@staticmethod` as a bound method, which is how the anchors reach the builder: the parent
    # calls `self.build_ref2va_packed_sequence(...)` positionally and knows nothing about keyframes.
    def build_ref2va_packed_sequence(self, *args, **kwargs):
        return build_combined_packed_sequence(*args, keyframe_anchors=self._keyframe_anchors, **kwargs)

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        self._keyframe_anchors = tuple(getattr(block_state, "keyframe_anchors", ()) or ())
        try:
            return super().__call__(components, state)
        finally:
            self._keyframe_anchors = ()


class MiniMaxH3CombinedSetupStep(SequentialPipelineBlocks):
    model_name = "minimax-h3"
    block_classes = [MiniMaxH3Ref2VASetupStep, MiniMaxH3KeyframesOnCanvasStep]
    block_names = ["references", "keyframes"]

    @property
    def description(self) -> str:
        return "Resolves the request plan from the references, then puts the keyframes on the resolved canvas."


class MiniMaxH3CombinedVaeEncoderStep(SequentialPipelineBlocks):
    model_name = "minimax-h3"
    block_classes = [MiniMaxH3Ref2VAReferenceEncoderStep, MiniMaxH3CombinedKeyframeEncoderStep]
    block_names = ["references", "keyframes"]

    @property
    def description(self) -> str:
        return "Encodes the references, then the keyframes, leaving the keyframe latents first in packed order."


class MiniMaxH3CombinedCoreDenoiseStep(SequentialPipelineBlocks):
    model_name = "minimax-h3"
    block_classes = [
        MiniMaxH3CombinedPrepareLayoutStep,
        MiniMaxH3PrepareConditionLatentsStep,
        MiniMaxH3PrepareLatentsStep,
        MiniMaxH3Ref2VAPrepareLatentsStep,
        MiniMaxH3SetTimestepsStep,
        MiniMaxH3Ref2VADenoiseStep,
        MiniMaxH3AfterDenoiseStep,
    ]
    block_names = [
        "prepare_layout",
        "prepare_condition_latents",
        "prepare_latents",
        "prepare_ref_latents",
        "set_timesteps",
        "denoise",
        "after_denoise",
    ]

    @property
    def description(self) -> str:
        return "Core denoising for a combined request: the `ref2va` chain over the combined packed layout."


def _is_combined(kwargs) -> bool:
    return kwargs.get("references") is not None and (
        kwargs.get("image") is not None or kwargs.get("last_image") is not None
    )


class MiniMaxH3CombinedAutoBeforeEncodeStep(MiniMaxH3AutoBeforeEncodeStep):
    block_classes = [MiniMaxH3CombinedSetupStep] + MiniMaxH3AutoBeforeEncodeStep.block_classes
    block_names = ["combined"] + MiniMaxH3AutoBeforeEncodeStep.block_names

    def select_block(self, **kwargs) -> str | None:
        return "combined" if _is_combined(kwargs) else super().select_block(**kwargs)


class MiniMaxH3CombinedAutoVaeEncoderStep(MiniMaxH3AutoVaeEncoderStep):
    block_classes = [MiniMaxH3CombinedVaeEncoderStep] + MiniMaxH3AutoVaeEncoderStep.block_classes
    block_names = ["combined"] + MiniMaxH3AutoVaeEncoderStep.block_names

    def select_block(self, **kwargs) -> str | None:
        return "combined" if _is_combined(kwargs) else super().select_block(**kwargs)


class MiniMaxH3CombinedAutoDenoiseStep(MiniMaxH3AutoDenoiseStep):
    block_classes = [MiniMaxH3CombinedCoreDenoiseStep] + MiniMaxH3AutoDenoiseStep.block_classes
    block_names = ["combined"] + MiniMaxH3AutoDenoiseStep.block_names

    def select_block(self, **kwargs) -> str | None:
        return "combined" if _is_combined(kwargs) else super().select_block(**kwargs)


class MiniMaxH3CombinedBlocks(MiniMaxH3Blocks):
    """`MiniMaxH3Blocks` plus a fourth shape: keyframes and references in the same generation.

    Supported workflows: `t2va`, `fl2va`, `ref2va` — unchanged — and `combined`, which needs `prompt`, `references`
    and at least one of `image` / `last_image`. A combined request runs against the `transformer_ref` partition, the
    same one `ref2va` uses.
    """

    block_classes = [
        MiniMaxH3CombinedAutoBeforeEncodeStep,
        MiniMaxH3AutoTextEncoderStep,
        MiniMaxH3CombinedAutoVaeEncoderStep,
        MiniMaxH3CombinedAutoDenoiseStep,
        MiniMaxH3DecodeStep,
    ]
    block_names = ["before_encode", "text_encoder", "vae_encoder", "denoise", "decode"]
    _workflow_map = dict(
        MiniMaxH3Blocks._workflow_map,
        combined=(
            {"prompt": True, "references": True, "image": True},
            {"prompt": True, "references": True, "last_image": True},
        ),
    )

    @property
    def description(self) -> str:
        return (
            "MiniMax-H3 blocks for joint video + audio generation, with the `t2va`, `fl2va` and `ref2va` workflows "
            "unchanged and a fourth, `combined`, that carries keyframe *and* reference conditioning in one packed "
            "sequence instead of dropping one of them."
        )