"""Caches the keyframe/condition-row slice of MiniMax-H3's video patch projection across denoising steps.

`MiniMaxH3Transformer3DModel.forward` projects the whole packed video latent through `proj_in` on every call,
including the leading condition (keyframe) rows -- which the denoise loop never re-steps (see
`MiniMaxH3LoopSchedulerStep`: `latents[num_condition_video_rows:] = scheduler.step(...)`, an in-place slice
assignment that never touches the condition rows at all). Those rows carry the same value, in the same storage,
for the whole denoise loop, so `proj_in`'s output for them is identical on every step -- recomputing it is pure
waste. Because `proj_in` is a plain linear layer, splitting it over the condition/generated row boundary and
concatenating is mathematically identical to running it over the whole tensor at once: this changes nothing about
the output, only when the work happens.

The cache key is the condition slice's own `data_ptr()` + shape + device, which is safe specifically because of the
in-place assignment above: the condition rows live in the same storage, unchanged, for an entire `pipe(...)` call,
and a *new* call (a new stage, a new chunk, an unrelated generation) always allocates a fresh latents tensor -- the
cache invalidates itself automatically, with no explicit per-request reset needed.

`patch()` wraps `forward` once, at load time, preserving the original's `@apply_lora_scale` behavior inline (a
real decorator -- it scales and unscales active LoRA layers around the forward pass, not just bookkeeping -- so
bypassing it would silently break LoRA strength). The wrapped method checks `transformer._proj_cache_enabled`
fresh on every call, so `set_enabled()` can toggle it per request without needing to unpatch anything.
"""

from __future__ import annotations

import torch

_PATCHED: set[int] = set()


def _condition_row_count(video_indices: torch.Tensor) -> int:
    """Number of leading, contiguous condition (keyframe) rows in `video_indices`' own order."""
    host = video_indices.detach().cpu()
    if host.numel() < 2:
        return 0
    discontinuities = (host[1:] - host[:-1] != 1).nonzero().flatten()
    return int(discontinuities[0]) + 1 if len(discontinuities) else 0


def patch(transformer) -> None:
    """Wrap `transformer.forward` once. Safe to call at startup; a no-op on a second call for the same instance."""
    if id(transformer) in _PATCHED:
        return

    from diffusers.models.modeling_utils import get_parameter_dtype
    from diffusers.models.transformers.transformer_minimax_h3 import (
        MINIMAX_H3_MODALITY_NUM,
        MiniMaxH3TransformerOutput,
    )
    from diffusers.utils import USE_PEFT_BACKEND
    from diffusers.utils.peft_utils import scale_lora_layers, unscale_lora_layers

    original_forward = transformer.forward

    def forward(
        hidden_states,
        audio_hidden_states,
        encoder_hidden_states,
        timestep,
        timestep_indices,
        token_tags,
        position_ids,
        video_indices,
        audio_indices,
        text_indices,
        attention_kwargs=None,
        return_dict=True,
    ):
        self = transformer
        if not getattr(self, "_proj_cache_enabled", True):
            return original_forward(
                hidden_states, audio_hidden_states, encoder_hidden_states, timestep, timestep_indices, token_tags,
                position_ids, video_indices, audio_indices, text_indices, attention_kwargs, return_dict,
            )

        # Verbatim from the stock `@apply_lora_scale("attention_kwargs")` decorator this replaces.
        lora_scale = 1.0
        kwargs = attention_kwargs
        if kwargs is not None:
            kwargs = kwargs.copy()
            lora_scale = kwargs.pop("scale", 1.0)
        if USE_PEFT_BACKEND:
            scale_lora_layers(self, lora_scale)
        try:
            if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
                raise ValueError(f"`position_ids` must be a `(seq_len, 3)` tensor, got {list(position_ids.shape)}.")
            sequence_length = position_ids.shape[0]
            if token_tags.shape != (sequence_length,) or timestep_indices.shape != (sequence_length,):
                raise ValueError(
                    "`token_tags` and `timestep_indices` must both be `(seq_len,)` tensors matching `position_ids`, "
                    f"got {list(token_tags.shape)} and {list(timestep_indices.shape)} for seq_len={sequence_length}."
                )

            rotary_emb = self.rope(position_ids)

            # The only change from the stock forward: `proj_in` is split at the condition/generated row boundary
            # and the condition half is cached. `torch.cat([proj_in(a), proj_in(b)]) == proj_in(torch.cat([a, b]))`
            # for a plain linear layer, so this is bit-identical to the stock computation, never an approximation.
            condition_rows = _condition_row_count(video_indices)
            if condition_rows > 0:
                dtype = get_parameter_dtype(self.proj_in)
                condition_slice = hidden_states[:, :condition_rows]
                key = (condition_slice.data_ptr(), tuple(condition_slice.shape), condition_slice.device)
                cached = getattr(self, "_condition_embeds_cache", None)
                if cached is not None and cached[0] == key:
                    condition_embeds = cached[1]
                else:
                    condition_embeds = self.proj_in(condition_slice.to(dtype))
                    self._condition_embeds_cache = (key, condition_embeds)
                generated_slice = hidden_states[:, condition_rows:]
                video_embeds = (
                    torch.cat((condition_embeds, self.proj_in(generated_slice.to(dtype))), dim=1)
                    if generated_slice.shape[1] > 0
                    else condition_embeds
                )
            else:
                video_embeds = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in)))

            audio_embeds = self.audio_proj_in(audio_hidden_states.to(get_parameter_dtype(self.audio_proj_in)))
            text_embeds = self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder)))
            text_embeds = self.token_refiner(text_embeds)

            hidden_states = text_embeds.new_zeros((text_embeds.shape[0], sequence_length, text_embeds.shape[-1]))
            hidden_states = hidden_states.index_copy(1, text_indices, text_embeds)
            hidden_states = hidden_states.index_copy(1, video_indices, video_embeds.to(text_embeds.dtype))
            hidden_states = hidden_states.index_copy(1, audio_indices, audio_embeds.to(text_embeds.dtype))

            temb = self.time_proj(timestep)
            temb = self.time_embedder(temb.to(get_parameter_dtype(self.time_embedder)))

            adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags

            for block in self.transformer_blocks:
                if torch.is_grad_enabled() and self.gradient_checkpointing:
                    hidden_states = self._gradient_checkpointing_func(
                        block, hidden_states, temb, adaln_indices, rotary_emb
                    )
                else:
                    hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb)

            hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(get_parameter_dtype(self.proj_out))
            video_output = self.proj_out(hidden_states).index_select(1, video_indices)
            audio_output = self.audio_proj_out(hidden_states).index_select(1, audio_indices)

            if not return_dict:
                return (video_output, audio_output)
            return MiniMaxH3TransformerOutput(sample=video_output, audio_sample=audio_output)
        finally:
            if USE_PEFT_BACKEND:
                unscale_lora_layers(self, lora_scale)

    transformer.forward = forward
    transformer._proj_cache_enabled = True
    _PATCHED.add(id(transformer))


def set_enabled(transformer, enabled: bool) -> None:
    """Toggle for one request. Call right before `pipe(...)`; `patch()` must have run first (a no-op otherwise,
    since an unpatched transformer just runs its own stock forward regardless of this attribute)."""
    transformer._proj_cache_enabled = bool(enabled)


def status() -> str:
    return "keyframe-projection cache patched" if _PATCHED else "keyframe-projection cache not patched"