"""Die Hälften eines **gesplitteten** MiniMax-H3-Deployments, plus eine Video-only-Variante.

MiniMax-H3 ist in bfloat16 195.9 GiB und ein ZeroGPU-Space wird bei 150 GB Storage evicted, also wird
`MiniMaxH3Blocks` an seinem `text_encoder`-Step aufgeschnitten: das 62.14 GiB Qwen3-VL laeuft im Conditioner-Space,
alles andere hier. `prompt_embeds` + `text_token_tags` ist das ganze Wire-Format dazwischen.

Zusaetzlich zu den Bloecken von akhaliq/MiniMax-H3-Turbo-Lora definiert dieses Modul
`MiniMaxH3MuteGeneratorBlocks`: identisch bis auf den Decode-Step, der nur `MiniMaxH3VideoDecodeStep` enthaelt.
Damit laesst sich pro Request entscheiden, ob der Audio-VAE ueberhaupt angefasst wird.

WICHTIG: Diese Klassen leiten von Block-Klassen des gepinnten diffusers-Commits ab (siehe requirements.txt).
Beim Neupinnen die Klassennamen unten gegen den neuen Head pruefen.
"""

from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
from diffusers.modular_pipelines.minimax_h3.decoders import (
    MiniMaxH3AfterDenoiseStep,
    MiniMaxH3VideoDecodeStep,
)
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]:
    """Das Wire-Format des Splits. `num_frames` deklariert nur die `ref2va`-Haelfte, deren Setup eines aufloest."""
    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):
    """Die Conditioner-Haelfte: Keyframes auf dem Canvas plus das Qwen3-VL, gelesen an 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):
    """Die Denoising-Haelfte: `MiniMaxH3Blocks` ohne seinen `text_encoder`-Step. Video **und** Soundtrack."""

    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 MiniMaxH3VideoOnlyDecodeStep(SequentialPipelineBlocks):
    """`MiniMaxH3DecodeStep` ohne den Audio-Zweig.

    Die Audio-Rows sind Teil derselben gepackten Sequenz und werden mitdenoised — das ist eine
    Checkpoint-Eigenschaft und laesst sich nicht abwaehlen. Was hier entfaellt, ist der Audio-VAE:
    er wird nicht deklariert, also auch nicht auf die Karte bewegt, nicht denormalisiert und nicht decodiert.
    Die denoisten Audio-Latents bleiben ungenutzt im PipelineState liegen.
    """

    model_name = "minimax-h3"
    block_classes = [MiniMaxH3VideoDecodeStep]
    block_names = ["video"]

    @property
    def description(self):
        return "Decodes only the video rows of the packed sequence; the audio VAE is never touched."

    @property
    def outputs(self):
        return [OutputParam.template("videos", description="The generated video.")]


class MiniMaxH3MuteGeneratorBlocks(SequentialPipelineBlocks):
    """`MiniMaxH3GeneratorBlocks` mit Video-only-Decode. Deklariert keinen `audio_vae`."""

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

    @property
    def description(self):
        return (
            "The denoising half of a split MiniMax-H3 deployment, video only: same blocks as "
            "`MiniMaxH3GeneratorBlocks` but with the audio decode step dropped, so the float32 audio VAE is neither "
            "declared nor moved onto the card."
        )

    @property
    def outputs(self):
        return [OutputParam.template("videos", description="The generated video.")]


class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
    """Die Conditioner-Haelfte eines gesplitteten `ref2va`. Komponentenweise identisch zu `MiniMaxH3ConditionerBlocks`."""

    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 and encodes "
            "MiniMax-H3's presentation of it into the `prompt_embeds` / `text_token_tags` pair."
        )

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


class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
    """Die Denoising-Haelfte eines gesplitteten `ref2va`."""

    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 transformer is the `transformer_ref` "
            "partition."
        )

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