# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

import functools
import inspect
import math
from dataclasses import dataclass
from typing import Any

import torch
from vllm.logger import init_logger
from vllm.model_executor.models.utils import extract_layer_index

from vllm_omni.diffusion.attention.backends.abstract import (
    AttentionBackend,
    AttentionImpl,
    AttentionMetadata,
)
from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available

logger = init_logger(__name__)

# The rf_v2 kernel only implements a 128-token block.
_BLOCK_SIZE = 128

# Below this many video blocks, the pooling and gather that block selection adds
# cost more than the QK work it removes, so stay dense.
_MIN_VIDEO_BLOCKS = 32

# rf_v2's ``input_layout`` describes the caller's tensors, and everything below
# slices the sequence on dim 1. vLLM-Omni diffusion attention hands the impl
# [B, S, N, D], so that is the only layout this backend accepts.
_INPUT_LAYOUT = "BSND"

# rf_v2 precision mode: 0 = high precision, 1 = high performance. Kept at the
# precise setting because the sparsity knob is the intended perf lever, and on
# A5 devices the kernel is routed to rf_v3, which overrides this anyway.
_INNER_PRECISE = 0

_WRONG_PLATFORM = (
    "RAINFUSION_ATTN runs the MindIE-SD rf_v2 kernel and is available on Ascend NPU only. "
    "Select FLASH_ATTN or TORCH_SDPA on this platform."
)

_MISSING_MINDIESD = (
    "RAINFUSION_ATTN requires MindIE-SD. Please install MindIE-SD to enable RainFusion sparse "
    "attention on Ascend NPU. For installation details, see https://gitcode.com/Ascend/MindIE-SD "
    "Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN"
)

_INCOMPATIBLE_MINDIESD = (
    "RAINFUSION_ATTN requires a MindIE-SD build whose sparse_attention supports the video_spans "
    "argument. Please upgrade MindIE-SD or select FLASH_ATTN."
)


# Whether the installed mindiesd ``sparse_attention`` accepts ``precision=``.
# Releases without it accept the kwarg through ``**kwargs`` but silently ignore
# it, so a requested mix/fp8 mode would silently run the BF16 path. Cached
# because forwards run per layer per denoise step.
@functools.cache
def _mindiesd_supports_precision() -> bool:
    try:
        from inspect import signature

        from mindiesd import sparse_attention

        return "precision" in signature(sparse_attention).parameters
    except Exception:
        return False


def _try_extract_layer_index(prefix: str) -> int | None:
    if not prefix:
        return None
    try:
        return extract_layer_index(prefix)
    except (AssertionError, ValueError):
        return None


def _supports_video_spans(sparse_attention: Any) -> bool:
    try:
        return "video_spans" in inspect.signature(sparse_attention).parameters
    except (TypeError, ValueError):
        return False


@dataclass(frozen=True)
class RainFusionConfig:
    """Resolved RainFusion controls for one attention layer.

    ``sparsity`` is the nominal fraction of key blocks dropped per query block.
    The realized sparsity is lower because rf_v2 always keeps the prefix rows and
    the first-frame blocks. ``start_step`` and ``skip_layers`` are the accuracy
    knobs: early denoise steps and specific DiT blocks stay dense.
    """

    sparsity: float = 0.0
    start_step: int = 0
    end_step: int = 0
    precision: str = "bf16"
    skip_layers: frozenset[int] = frozenset()

    @classmethod
    def from_backend_kwargs(cls, backend_kwargs: dict | None) -> RainFusionConfig:
        bk = backend_kwargs or {}
        return cls(
            sparsity=float(bk.get("sparsity", 0.0)),
            start_step=int(bk.get("start_step", 0)),
            end_step=int(bk.get("end_step", 0)),
            precision=str(bk.get("precision", "bf16")),
            skip_layers=frozenset(bk.get("skip_layers") or ()),
        )

    @property
    def enabled(self) -> bool:
        return self.sparsity > 0.0


@dataclass(frozen=True)
class RainFusionPlan:
    """Per-forward geometry handed to the rf_v2 kernel."""

    used_len: int
    prefix_len: int | None = None
    latent_shape: list[int] | None = None
    video_spans: list[dict[str, object]] | None = None


class RainFusionAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    supported_platforms: tuple[str, ...] = ("npu",)

    @classmethod
    def validate_available(cls) -> None:
        from importlib.util import find_spec

        if find_spec("mindiesd") is None:
            raise ValueError(_MISSING_MINDIESD)
        try:
            from mindiesd import sparse_attention
        except ImportError as exc:
            raise ValueError(_MISSING_MINDIESD) from exc
        if not _supports_video_spans(sparse_attention):
            raise ValueError(_INCOMPATIBLE_MINDIESD)

    @staticmethod
    def get_supported_head_sizes() -> list[int]:
        return [64, 96, 128, 192, 256]

    @staticmethod
    def get_name() -> str:
        return "RAINFUSION_ATTN"

    @staticmethod
    def get_impl_cls() -> type[RainFusionAttentionImpl]:
        return RainFusionAttentionImpl


class RainFusionAttentionImpl(AttentionImpl):
    """Block-sparse video attention via MindIE-SD RainFusion (rf_v2) on Ascend NPU.

    Sparsity applies only to the video segment of a packed multimodal sequence,
    whose extent the model publishes as ``AttentionMetadata.video_layout``. Every
    other case — warmup denoise steps, exempt layers, a layer that does not declare
    ``qkv_layout="BSND"``, sequences without a published video segment, video
    segments too short to pay for block selection — delegates to FlashAttention,
    so a model can select this backend unconditionally. MindIE-SD handles an
    irregular video tail internally, retaining it outside the sparse blocks.
    """

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
        num_kv_heads: int | None = None,
        prefix: str = "",
        qkv_layout: str | None = None,
        backend_kwargs: dict[str, Any] | None = None,
        **extra_impl_args,
    ) -> None:
        self.num_heads = num_heads
        self.causal = causal
        self.softmax_scale = softmax_scale
        self.qkv_layout = qkv_layout

        self.rainfusion = RainFusionConfig.from_backend_kwargs(backend_kwargs)
        self.layer_idx = _try_extract_layer_index(prefix)

        if self.rainfusion.enabled:
            self._validate_parallel_config()
            if causal:
                raise ValueError(
                    "RAINFUSION_ATTN does not support causal attention: rf_v2 selects key "
                    "blocks by pooled relevance and cannot express a causal mask. Select "
                    "FLASH_ATTN for causal roles."
                )
            if qkv_layout is not None and qkv_layout.upper() != _INPUT_LAYOUT:
                raise ValueError(
                    f"RAINFUSION_ATTN needs {_INPUT_LAYOUT} tensors to locate the video segment along "
                    f"the sequence axis, but this layer declares qkv_layout={qkv_layout!r}. Select "
                    "FLASH_ATTN for this role."
                )

        self.dense_fallback = FlashAttentionBackend.get_impl_cls()(
            num_heads=num_heads,
            head_size=head_size,
            softmax_scale=softmax_scale,
            causal=causal,
            num_kv_heads=num_kv_heads,
            prefix=prefix,
            qkv_layout=qkv_layout,
        )

    def _validate_parallel_config(self) -> None:
        config = get_current_diffusion_config_or_none()
        parallel_config = getattr(config, "parallel_config", None)
        ring_degree = getattr(parallel_config, "ring_degree", 1)
        if ring_degree > 1:
            # Ring gives each rank a slice of the sequence, so block selection
            # would score only local keys and the layer bypasses the backend
            # entirely (see Attention._run_ring_attention).
            raise ValueError(
                "RAINFUSION_ATTN is not compatible with ring sequence parallelism "
                f"(ring_degree={ring_degree}): rf_v2 needs the whole key sequence to rank "
                "blocks. Use Ulysses SP (ring_degree=1) instead."
            )

    def forward_cuda(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        # ROCm and MUSA route through forward_cuda by default, so this covers them too.
        raise NotImplementedError(_WRONG_PLATFORM)

    def forward_xpu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError(_WRONG_PLATFORM)

    def forward_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        plan = self._resolve_plan(attn_metadata)
        if plan is None:
            return self.dense_fallback.forward_npu(query, key, value, attn_metadata)
        return self._forward_sparse_npu(query, key, value, plan)

    def _resolve_plan(self, attn_metadata: AttentionMetadata | None) -> RainFusionPlan | None:
        """Return the rf_v2 geometry, or None when this forward must stay dense."""
        rf = self.rainfusion
        if not rf.enabled:
            return None
        if self.layer_idx is not None and self.layer_idx in rf.skip_layers:
            return None
        if is_forward_context_available():
            step_idx = get_forward_context().denoise_step_idx
            total_steps = get_forward_context().total_denoise_steps
            if step_idx is not None and step_idx < rf.start_step:
                return None
            # Tail fallback: keep the last ``end_step`` denoise steps dense.
            if (
                rf.end_step > 0
                and step_idx is not None
                and total_steps is not None
                and step_idx >= total_steps - rf.end_step
            ):
                return None
        if self.qkv_layout is None:
            # The sparse path reads the sequence off dim 1, which the tensors alone
            # do not establish, and the dense fallback resolves an absent layout its
            # own way. Sparsifying on an assumption would put the two paths on
            # different axes, so an undeclared layout stays dense.
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: this layer does not declare qkv_layout, and rf_v2 "
                "needs %s to locate the video segment along the sequence axis. Set qkv_layout=%r on "
                "the Attention layer to enable sparsity.",
                _INPUT_LAYOUT,
                _INPUT_LAYOUT,
            )
            return None

        if attn_metadata is None:
            return None

        layout = attn_metadata.video_layout
        if layout is None:
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: this attention role carries no video segment. The "
                "model must publish AttentionMetadata.video_layout for the sequence to be sparsified."
            )
            return None
        max_seqlen_q = attn_metadata.extra.get("max_seqlen_q")
        if max_seqlen_q is None:
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: attention metadata is missing max_seqlen_q, so the "
                "video segment cannot be confirmed to be the tail of packed document 0."
            )
            return None

        if layout.video_spans:
            return self._resolve_multi_span_plan(layout, max_seqlen_q)

        if layout.prefix_len is None or layout.latent_grid is None:
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: video layout has neither a legacy video tail nor multi-video spans."
            )
            return None
        prefix_len = int(layout.prefix_len)
        latent_shape = [int(dim) for dim in layout.latent_grid]
        # rf_v2 splits the sequence as [prefix | t*h*w video rows]. Document 0 of
        # the packed sequence holds those rows; anything past it is alignment
        # padding that rf_v2 must not see.
        video_len = math.prod(latent_shape)
        used_len = prefix_len + video_len

        if used_len != int(max_seqlen_q):
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: prefix (%d) plus latent grid %s does not fill "
                "packed document 0 (%d rows). rf_v2 requires the video segment to be its tail.",
                prefix_len,
                tuple(latent_shape),
                int(max_seqlen_q),
            )
            return None
        if video_len < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE:
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: %d video rows is under the %d-row "
                "(%d block) threshold where sparse selection pays off.",
                video_len,
                _MIN_VIDEO_BLOCKS * _BLOCK_SIZE,
                _MIN_VIDEO_BLOCKS,
            )
            return None
        logger.info_once(
            "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, "
            "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "
            "because prefix and first-frame blocks are always kept.",
            rf.sparsity,
            rf.start_step,
            len(rf.skip_layers),
            tuple(latent_shape),
            prefix_len,
            video_len,
        )
        return RainFusionPlan(
            used_len=used_len,
            prefix_len=prefix_len,
            latent_shape=latent_shape,
        )

    def _resolve_multi_span_plan(self, layout, max_seqlen_q: int) -> RainFusionPlan | None:
        """Validate Ref2VA's non-contiguous video grids before sparse dispatch."""
        used_len = layout.used_len
        if used_len is None or int(used_len) != int(max_seqlen_q):
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: multi-video layout used_len=%r does not match packed document 0 (%d).",
                used_len,
                int(max_seqlen_q),
            )
            return None

        spans: list[dict[str, object]] = []
        span_summaries: list[str] = []
        previous_end = 0
        target_count = 0
        video_seqlen = 0
        boundary_dense_seqlen = 0
        previous_video_length: int | None = None
        for span in sorted(layout.video_spans, key=lambda item: item.start):
            grid = tuple(int(dim) for dim in span.latent_grid)
            length = math.prod(grid)
            start = int(span.start)
            if (
                len(grid) != 3
                or any(dim <= 0 for dim in grid)
                or start < previous_end
                or start + length > int(used_len)
            ):
                logger.warning_once(
                    "RAINFUSION_ATTN staying dense: invalid multi-video span start=%d grid=%s used_len=%d.",
                    start,
                    grid,
                    int(used_len),
                )
                return None
            role = span.role
            if role not in ("reference", "target"):
                logger.warning_once("RAINFUSION_ATTN staying dense: unsupported multi-video span role %r.", role)
                return None
            if role == "target":
                target_count += 1
            if previous_video_length is not None:
                # rf_v2 works on fixed 128-token blocks. The preceding clip
                # needs these dense rows to complete its tail block before
                # this clip begins, otherwise one sparse block would cross a
                # clip boundary.
                boundary_dense_seqlen += (-previous_video_length) % _BLOCK_SIZE
            spans.append({"start": start, "latent_shape": list(grid)})
            span_summaries.append(f"role={role}, start={start}, seqlen={length}, latent_shape={grid}")
            previous_end = start + length
            video_seqlen += length
            previous_video_length = length

        if target_count != 1:
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: Ref2VA layout must contain exactly one target video span, got %d.",
                target_count,
            )
            return None
        if video_seqlen < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE:
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: multi-video seqlen=%d is under the sparse threshold seqlen=%d.",
                video_seqlen,
                _MIN_VIDEO_BLOCKS * _BLOCK_SIZE,
            )
            return None
        dense_context_seqlen = int(used_len) - video_seqlen
        if boundary_dense_seqlen > dense_context_seqlen:
            logger.warning_once(
                "RAINFUSION_ATTN staying dense: multi-video spans need dense_context_seqlen=%d "
                "to isolate clip block boundaries, but this layout has only %d.",
                boundary_dense_seqlen,
                dense_context_seqlen,
            )
            return None
        logger.info_once(
            "RAINFUSION_ATTN multi-video active: sparsity=%.2f, spans=[%s], "
            "video_seqlen=%d, dense_context_seqlen=%d, valid_packed_seqlen=%d.",
            self.rainfusion.sparsity,
            "; ".join(span_summaries),
            video_seqlen,
            dense_context_seqlen,
            int(used_len),
        )
        return RainFusionPlan(used_len=int(used_len), video_spans=spans)

    def _forward_sparse_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        plan: RainFusionPlan,
    ) -> torch.Tensor:
        try:
            from mindiesd import sparse_attention
        except ImportError:
            raise ImportError(_MISSING_MINDIESD)
        if self.rainfusion.precision != "bf16" and not _mindiesd_supports_precision():
            raise RuntimeError(
                f"block_sparse.precision={self.rainfusion.precision!r} requires MindIE-SD "
                "with sparse_attention(precision=...) support; the installed mindiesd "
                "silently ignores it and would run the BF16 path. Install a compatible "
                "MindIE-SD release or use precision='bf16'."
            )

        used = plan.used_len
        q, k, v = (tensor[:, :used] for tensor in (query, key, value))
        # Ulysses has already gathered the full sequence onto this rank and split
        # the heads, so read the head count off the tensor rather than num_heads.
        common_kwargs: dict[str, object] = {
            "scale": self.softmax_scale,
            "head_num": query.shape[-2],
            "input_layout": _INPUT_LAYOUT,
            "inner_precise": _INNER_PRECISE,
            "block_size": _BLOCK_SIZE,
            "sparsity": self.rainfusion.sparsity,
            "precision": self.rainfusion.precision,
        }
        if plan.video_spans is not None:
            out = sparse_attention(
                q,
                k,
                v,
                sparse_type="rf_v2",
                video_spans=plan.video_spans,
                **common_kwargs,
            )
        else:
            assert plan.prefix_len is not None and plan.latent_shape is not None
            common_kwargs.update(
                sparse_type="rf_v2",
                txt_len=plan.prefix_len,
                latent_shape_q=plan.latent_shape,
                latent_shape_k=plan.latent_shape,
            )
            out = sparse_attention(q, k, v, **common_kwargs)
        if used == query.shape[1]:
            return out
        padded = torch.zeros_like(query)
        padded[:, :used] = out
        return padded