"""The halves of a **split** MiniMax-H3 deployment, for both of its checkpoint partitions.

MiniMax-H3 is modular-only, and the whole model is one `MiniMaxH3Blocks` sequence whose branches are picked per
request — and per `workflow=` — from the inputs:

    before_encode -> text_encoder -> vae_encoder -> denoise -> after_denoise -> decode

where `before_encode`, `text_encoder`, `vae_encoder` and `denoise` each switch on `references` (the `ref2va` workflow)
versus the keyframe inputs (`t2va` / `fl2va`), and `denoise` is itself `prepare_layout -> prepare_latents ->
set_timesteps -> denoise` against `transformer` or `transformer_ref`.

The conditioner (a 62.14 GiB Qwen3-VL) and the denoiser (a 61.73 GiB transformer plus ~20.5 GiB of float32 VAEs) do
not fit on one 95 GiB card unquantized, so this module cuts that sequence in two at the `text_encoder` step, once per
partition:

  * `MiniMaxH3ConditionerBlocks` = `[resize, text_encoder]` — loads `text_encoder` / `tokenizer` / `processor` only
    (plus the `image_processor`, which is built from config and downloads nothing), and emits `prompt_embeds` +
    `text_token_tags`, which is the whole wire format between the two halves.
  * `MiniMaxH3GeneratorBlocks` = everything else — loads `transformer` / `vae` / `audio_vae` / the two schedulers
    only, and takes `prompt_embeds` + `text_token_tags` as *inputs*.
  * `MiniMaxH3Ref2VAConditionerBlocks` / `MiniMaxH3Ref2VAGeneratorBlocks` are the same cut through the `ref2va`
    branch, so one conditioner Space serves both partitions out of the weights it already holds.

`resize` / `setup` run on both sides on purpose. They own no pretrained component (PIL, decoded media and arithmetic),
they resolve the canvas and prepare the keyframes or normalize the references — which the conditioner needs to build
its vision blocks and the generator needs to encode with the VAEs. Running them twice over the same inputs is
deterministic; both conditioner halves return the resolved `height` / `width` / `num_frames` anyway, so the caller
pins them explicitly on the generating half.

Two things the blocks leave to the caller: a keyframe reaches them EXIF-transposed and in RGB, and the `t2va` /
`fl2va` frame count is aligned to `17 * n + 5` before the call, since that arithmetic lives in the layout step on the
denoising side of the cut. `ref2va` still resolves its own frame count, but requires one to be passed.

Only *text* encoding is remote. `vae_encoder` / `reference_encoder` stay on the denoising side: they run the two
autoencoders, which the conditioner Space does not hold.
"""

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.encoders import (
    MiniMaxH3Ref2VAReferenceEncoderStep,
    MiniMaxH3Ref2VATextEncoderStep,
    MiniMaxH3TextEncoderStep,
)
from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
    MiniMaxH3AutoKeyframeVaeEncoderStep,
    MiniMaxH3AutoResizeStep,
    MiniMaxH3CoreDenoiseStep,
    MiniMaxH3DecodeStep,
    MiniMaxH3Ref2VACoreDenoiseStep,
    _generation_outputs,
)
from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam


def _wire_outputs(num_frames: bool = True) -> list[OutputParam]:
    """The wire format of the split, plus the plan the caller pins on the generating half.

    `num_frames` is declared by the `ref2va` half alone: it is the one whose setup step resolves a frame count.
    """
    return [
        OutputParam.template("prompt_embeds"),
        OutputParam("text_token_tags", description="The per-row modality tag of every row of `prompt_embeds`."),
        OutputParam("height", type_hint=int, description="Resolved height of the generated video in pixels."),
        OutputParam("width", type_hint=int, description="Resolved width of the generated video in pixels."),
        *(
            [OutputParam("num_frames", type_hint=int, description="Resolved number of frames, of the form 17 * n + 5.")]
            if num_frames
            else []
        ),
    ]


class MiniMaxH3ConditionerBlocks(SequentialPipelineBlocks):
    """The conditioner half of a split MiniMax-H3: the keyframes on the canvas plus the Qwen3-VL read at layer 50."""

    model_name = "minimax-h3"
    block_classes = [MiniMaxH3AutoResizeStep, MiniMaxH3TextEncoderStep]
    block_names = ["resize", "text_encoder"]

    @property
    def description(self):
        return (
            "The conditioner half of a split MiniMax-H3 deployment: puts the keyframes onto the target canvas and "
            "encodes MiniMax-H3's presentation of the request into the `prompt_embeds` / `text_token_tags` pair the "
            "denoising half consumes. The frame count is the caller's to align."
        )

    @property
    def outputs(self):
        return _wire_outputs(num_frames=False)


class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
    """The denoising half of a split MiniMax-H3: `MiniMaxH3Blocks` with its `text_encoder` step removed."""

    model_name = "minimax-h3"
    block_classes = [
        MiniMaxH3AutoResizeStep,
        MiniMaxH3AutoKeyframeVaeEncoderStep,
        MiniMaxH3CoreDenoiseStep,
        MiniMaxH3AfterDenoiseStep,
        MiniMaxH3DecodeStep,
    ]
    block_names = ["resize", "vae_encoder", "denoise", "after_denoise", "decode"]

    @property
    def description(self):
        return (
            "The denoising half of a split MiniMax-H3 deployment: the `t2va` / `fl2va` branch of `MiniMaxH3Blocks` "
            "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
            "62.14 GiB Qwen3-VL conditioner is never loaded here."
        )

    @property
    def outputs(self):
        return _generation_outputs()


class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
    """The conditioner half of a split `ref2va`: the resolved plan plus the Qwen3-VL read at its 50th layer.

    Component for component this is `MiniMaxH3ConditionerBlocks` — `text_encoder`, `tokenizer`, `processor` — which
    is what lets one conditioner Space serve both partitions of the checkpoint out of the weights it already holds.
    What differs is the presentation the Qwen3-VL is shown: `ref2va` prepends a label per reference, numbered per
    modality, and a vision block per image and per merged video frame pair, so the references themselves have to
    reach this half. An audio reference never does — it contributes its `"<Audio j>: "` label and nothing else — but
    it is still part of the request here, because the setup step normalizes every soundtrack and validates the mix.
    """

    model_name = "minimax-h3"
    block_classes = [MiniMaxH3Ref2VASetupStep, MiniMaxH3Ref2VATextEncoderStep]
    block_names = ["setup", "text_encoder"]

    @property
    def description(self):
        return (
            "The conditioner half of a split MiniMax-H3 `ref2va` deployment: resolves the request plan (canvas, frame "
            "count, references normalized onto MiniMax-H3's own rates and resolutions) and encodes MiniMax-H3's "
            "presentation of it into the `prompt_embeds` / `text_token_tags` pair the denoising half consumes."
        )

    @property
    def outputs(self):
        return _wire_outputs()


class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
    """The denoising half of a split `ref2va`: the `ref2va` branch with its `text_encoder` step removed.

    Only the text-encoder step is dropped. `reference_encoder` is this half's own encoder — it runs the video VAE
    over the image and video references and the audio VAE over the soundtracks, and its output shapes are where every
    reference block's geometry in the packed layout comes from — so it stays here, next to the autoencoders.
    """

    model_name = "minimax-h3"
    block_classes = [
        MiniMaxH3Ref2VASetupStep,
        MiniMaxH3Ref2VAReferenceEncoderStep,
        MiniMaxH3Ref2VACoreDenoiseStep,
        MiniMaxH3AfterDenoiseStep,
        MiniMaxH3DecodeStep,
    ]
    block_names = ["setup", "reference_encoder", "denoise", "after_denoise", "decode"]

    @property
    def description(self):
        return (
            "The denoising half of a split MiniMax-H3 `ref2va` deployment: the `ref2va` branch of `MiniMaxH3Blocks` "
            "without its text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the "
            "62.14 GiB Qwen3-VL conditioner is never loaded here. The transformer is the `transformer_ref` partition."
        )

    @property
    def outputs(self):
        return _generation_outputs()