# Copyright 2026 The MiniMax Team and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from ...utils import logging
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import InputParam, OutputParam
from .modular_pipeline import MiniMaxMusic3ModularPipeline


logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

# Hidden-state chunking: the autoregressive frames are decoded in 200-frame windows with a 100-frame hop; neighboring
# windows overlap by ~344 latent frames (~3.445 latents per frame), of which the trailing 86 latent frames
# (86 * 512 samples) are kept from the previous window when cropping the decoded waveform.
_CHUNK_FRAMES = 200
_CHUNK_HOP = 100


class MiniMaxMusic3PrepareChunksStep(ModularPipelineBlocks):
    model_name = "minimax-music3"

    @property
    def description(self) -> str:
        return (
            "Chunk bookkeeping step that splits the autoregressive frames into 200-frame windows with a 100-frame "
            "hop; each window is flow-matched with the previous window's trailing latents as an overlap prompt."
        )

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "frame_hiddens",
                required=True,
                type_hint=torch.Tensor,
                description="Per-frame hidden states generated by the autoregressive step.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "chunk_starts",
                type_hint=list,
                description="Frame index at which each 200-frame denoising window starts.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: MiniMaxMusic3ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        num_frames = block_state.frame_hiddens.shape[1]
        block_state.chunk_starts = (
            [0] if num_frames <= _CHUNK_FRAMES else list(range(0, num_frames - _CHUNK_HOP, _CHUNK_HOP))
        )

        self.set_block_state(state, block_state)
        return components, state