# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Guarded Wan VAE decoder data-movement fast paths.

This module adapts the lossless Wan decoder optimizations from SGLang while
keeping Diffusers' parameter names and module topology unchanged. Every
optimized method falls back to the Diffusers 0.38.0 operation order when the
kernel, device, layout, dtype, autograd state, or module type is unsupported.
"""

from __future__ import annotations

from types import MethodType
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.models.autoencoders.autoencoder_kl_wan import (
    CACHE_T,
    DupUp3D,
    WanCausalConv3d,
    WanDecoder3d,
    WanResample,
    WanResidualBlock,
    WanResidualUpBlock,
)

from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard

try:
    from vllm_omni.diffusion.kernels.wan_vae_data_movement import (
        cat_pad_5d,
        dup_up3d_add,
    )
except ImportError:  # pragma: no cover - Triton is optional on non-GPU installs.
    cat_pad_5d = None
    dup_up3d_add = None

_INSTALL_ATTR = "_vllm_omni_wan_data_movement_installed"
_DYNAMIC_SPATIAL_CONV_ATTR = "_vllm_omni_dynamic_spatial_shard_conv"


def _cache_payload(cache: Any) -> torch.Tensor | None:
    return cache if isinstance(cache, torch.Tensor) else None


def _is_dynamic_spatial_conv(conv: nn.Module) -> bool:
    return bool(getattr(conv, _DYNAMIC_SPATIAL_CONV_ATTR, False))


def _spatial_shard_is_active() -> bool:
    return wan_spatial_shard._SPATIAL_SHARD_CONTEXT.get() is not None


def _causal_padding(conv: nn.Module) -> tuple[int, ...]:
    if _is_dynamic_spatial_conv(conv):
        return tuple(conv._source_padding)
    return tuple(conv._padding)


def _can_use_fused_cache_path(conv: nn.Module, x: torch.Tensor) -> bool:
    return bool(
        cat_pad_5d is not None
        and (type(conv) is WanCausalConv3d or _is_dynamic_spatial_conv(conv))
        and x.dim() == 5
        and x.is_cuda
        and not torch.is_grad_enabled()
        and not torch.compiler.is_compiling()
        and not _spatial_shard_is_active()
    )


def _run_cached_causal_conv(
    conv: nn.Module,
    x: torch.Tensor,
    cache_list: list[Any],
    index: int,
) -> torch.Tensor:
    """Run a causal conv and refresh its feature-cache slot."""
    cache = cache_list[index]
    is_repeat_marker = isinstance(cache, str)
    payload = None if is_repeat_marker else _cache_payload(cache)
    if _can_use_fused_cache_path(conv, x) and (
        payload is None or (payload.device == x.device and payload.dtype == x.dtype)
    ):
        pair = cat_pad_5d(x, payload, _causal_padding(conv), keep_cache_frames=CACHE_T)
        if pair is not None:
            conv_input, cache_list[index] = pair
            return nn.Conv3d.forward(conv, conv_input)

    cache_x = x[:, :, -CACHE_T:, :, :].clone()
    if cache_x.shape[2] < CACHE_T and payload is not None:
        cache_x = torch.cat(
            [payload[:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
            dim=2,
        )
    elif cache_x.shape[2] < CACHE_T and is_repeat_marker:
        cache_x = torch.cat([torch.zeros_like(cache_x), cache_x], dim=2)

    out = conv(x) if payload is None else conv(x, payload)
    cache_list[index] = cache_x
    return out


def _causal_conv_forward(self: WanCausalConv3d, x: torch.Tensor, cache_x: torch.Tensor | None = None):
    padding = list(self._padding)
    if (
        any(padding)
        and _can_use_fused_cache_path(self, x)
        and (cache_x is None or (cache_x.device == x.device and cache_x.dtype == x.dtype))
    ):
        conv_input = cat_pad_5d(x, cache_x, padding)
        if conv_input is not None:
            return nn.Conv3d.forward(self, conv_input)

    if cache_x is not None and padding[4] > 0:
        cache_x = cache_x.to(x.device)
        x = torch.cat([cache_x, x], dim=2)
        padding[4] -= cache_x.shape[2]
    x = F.pad(x, padding)
    return nn.Conv3d.forward(self, x)


def _resample_forward(
    self: WanResample,
    x: torch.Tensor,
    feat_cache: list[Any] | None = None,
    feat_idx: list[int] | None = None,
):
    if feat_idx is None:
        feat_idx = [0]
    batch, channels, frames, height, width = x.size()
    if self.mode == "upsample3d" and feat_cache is not None:
        index = feat_idx[0]
        if feat_cache[index] is None:
            feat_cache[index] = "Rep"
            feat_idx[0] += 1
        else:
            x = _run_cached_causal_conv(self.time_conv, x, feat_cache, index)
            feat_idx[0] += 1
            x = x.reshape(batch, 2, channels, frames, height, width)
            x = torch.stack((x[:, 0], x[:, 1]), 3)
            x = x.reshape(batch, channels, frames * 2, height, width)

    frames = x.shape[2]
    x = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width)
    x = self.resample(x)
    x = x.view(batch, frames, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4)

    if self.mode == "downsample3d" and feat_cache is not None:
        index = feat_idx[0]
        if feat_cache[index] is None:
            feat_cache[index] = x.clone()
            feat_idx[0] += 1
        else:
            cache_x = x[:, :, -1:, :, :].clone()
            x = self.time_conv(torch.cat([feat_cache[index][:, :, -1:, :, :], x], dim=2))
            feat_cache[index] = cache_x
            feat_idx[0] += 1
    return x


def _residual_block_forward(
    self: WanResidualBlock,
    x: torch.Tensor,
    feat_cache: list[Any] | None = None,
    feat_idx: list[int] | None = None,
):
    if feat_idx is None:
        feat_idx = [0]
    residual = self.conv_shortcut(x)
    x = self.nonlinearity(self.norm1(x))
    if feat_cache is not None:
        index = feat_idx[0]
        x = _run_cached_causal_conv(self.conv1, x, feat_cache, index)
        feat_idx[0] += 1
    else:
        x = self.conv1(x)

    x = self.dropout(self.nonlinearity(self.norm2(x)))
    if feat_cache is not None:
        index = feat_idx[0]
        x = _run_cached_causal_conv(self.conv2, x, feat_cache, index)
        feat_idx[0] += 1
    else:
        x = self.conv2(x)
    return x + residual


def _residual_up_block_forward(
    self: WanResidualUpBlock,
    x: torch.Tensor,
    feat_cache: list[Any] | None = None,
    feat_idx: list[int] | None = None,
    first_chunk: bool = False,
):
    if feat_idx is None:
        feat_idx = [0]
    shortcut_source = x.clone()
    for resnet in self.resnets:
        if feat_cache is None:
            x = resnet(x)
        else:
            x = resnet(x, feat_cache=feat_cache, feat_idx=feat_idx)

    if self.upsampler is not None:
        if feat_cache is None:
            x = self.upsampler(x)
        else:
            x = self.upsampler(x, feat_cache=feat_cache, feat_idx=feat_idx)

    shortcut = self.avg_shortcut
    if shortcut is None:
        return x
    if (
        dup_up3d_add is not None
        and type(shortcut) is DupUp3D
        and x.is_cuda
        and shortcut_source.is_cuda
        and x.dtype == shortcut_source.dtype
        and not torch.is_grad_enabled()
        and not torch.compiler.is_compiling()
        and not _spatial_shard_is_active()
    ):
        fused = dup_up3d_add(
            x,
            shortcut_source,
            shortcut.factor_t,
            shortcut.factor_s,
            shortcut.repeats,
            first_chunk,
        )
        if fused is not None:
            return fused
    return x + shortcut(shortcut_source, first_chunk=first_chunk)


def _decoder_forward(
    self: WanDecoder3d,
    x: torch.Tensor,
    feat_cache: list[Any] | None = None,
    feat_idx: list[int] | None = None,
    first_chunk: bool = False,
):
    if feat_idx is None:
        feat_idx = [0]
    if feat_cache is not None:
        index = feat_idx[0]
        x = _run_cached_causal_conv(self.conv_in, x, feat_cache, index)
        feat_idx[0] += 1
    else:
        x = self.conv_in(x)

    x = self.mid_block(x, feat_cache=feat_cache, feat_idx=feat_idx)
    for up_block in self.up_blocks:
        x = up_block(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk)

    x = self.nonlinearity(self.norm_out(x))
    if feat_cache is not None:
        index = feat_idx[0]
        x = _run_cached_causal_conv(self.conv_out, x, feat_cache, index)
        feat_idx[0] += 1
    else:
        x = self.conv_out(x)
    return x


def install_wan_vae_data_movement(vae: nn.Module) -> bool:
    """Install guarded fast paths on a Diffusers Wan decoder once."""
    if getattr(vae, _INSTALL_ATTR, False):
        return True
    if cat_pad_5d is None or dup_up3d_add is None:
        return False

    decoder = getattr(vae, "decoder", None)
    if type(decoder) is not WanDecoder3d:
        return False

    decoder.forward = MethodType(_decoder_forward, decoder)
    patched = 0
    for module in decoder.modules():
        module_type = type(module)
        if module_type is WanCausalConv3d:
            module.forward = MethodType(_causal_conv_forward, module)
            patched += 1
        elif _is_dynamic_spatial_conv(module):
            # The dynamic wrapper owns halo exchange while a spatial request is
            # active. Its direct path can still use the lossless fused cache
            # movement through ``_run_cached_causal_conv``.
            patched += 1
        elif module_type is WanResidualBlock:
            module.forward = MethodType(_residual_block_forward, module)
        elif module_type is WanResidualUpBlock:
            module.forward = MethodType(_residual_up_block_forward, module)
        elif module_type is WanResample:
            module.forward = MethodType(_resample_forward, module)

    if patched == 0:
        return False
    setattr(vae, _INSTALL_ATTR, True)
    return True