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

from __future__ import annotations

import sys
from dataclasses import replace
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch

from vllm_omni.diffusion.attention.backends.abstract import (
    AttentionBackend,
    AttentionMetadata,
)
from vllm_omni.diffusion.attention.backends.flash_attn import (
    FlashAttentionBackend,
    FlashAttentionImpl,
)
from vllm_omni.diffusion.attention.backends.utils import fa
from vllm_omni.diffusion.attention.capabilities import (
    CapabilityResult,
    CompilationMode,
    ExecutionContext,
    OuterBoundary,
    ParallelStrategy,
    SupportStatus,
)

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

_validate_fa4_head_dims = fa.validate_fa4_head_dims


@pytest.fixture(autouse=True)
def kernel_validator(monkeypatch):
    # Capability unit tests do not require FA4 or a CUDA device. Real-kernel
    # validation is exercised by test_flash_attn_compile.py.
    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True)


def _impl(*, kernel_variant: str | None = "fa4", causal: bool = False):
    impl = FlashAttentionImpl.__new__(FlashAttentionImpl)
    impl._kernel_variant = kernel_variant
    impl.causal = causal
    return impl


def _resolve(
    impl,
    *,
    context: ExecutionContext | None = None,
    attn_metadata: AttentionMetadata | None = None,
    dtype=torch.bfloat16,
):
    tensor = torch.empty((1, 16, 8, 64), dtype=dtype)
    return impl.resolve_execution_path(
        context or ExecutionContext(platform="cuda"),
        tensor,
        tensor,
        tensor,
        attn_metadata,
    )


class _UnmigratedBackend(AttentionBackend):
    @staticmethod
    def get_name() -> str:
        return "UNMIGRATED_TEST"

    @staticmethod
    def get_impl_cls():
        return None

    @staticmethod
    def get_metadata_cls():
        return None

    @staticmethod
    def get_builder_cls():
        return None

    @staticmethod
    def get_supported_head_sizes() -> list[int]:
        return []


def test_capability_result_is_tri_state():
    assert CapabilityResult.supported().status is SupportStatus.SUPPORTED
    assert CapabilityResult.unsupported("not available").status is SupportStatus.UNSUPPORTED
    assert CapabilityResult.unmigrated().status is SupportStatus.UNMIGRATED
    with pytest.raises(ValueError, match="actionable reason"):
        CapabilityResult(SupportStatus.UNSUPPORTED)


def test_unmigrated_backend_default_is_conservative():
    context = ExecutionContext(platform="cuda", require_fullgraph=True)
    result = _UnmigratedBackend.resolve_capabilities(context)

    assert result.support.status is SupportStatus.UNMIGRATED
    assert result.compilation_mode is CompilationMode.EAGER_ONLY
    requested = result.requested_support(context)
    assert requested.status is SupportStatus.UNSUPPORTED
    assert "no verified fullgraph declaration" in requested.reason
    assert result.support.status is SupportStatus.UNMIGRATED
    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.UNMIGRATED


@pytest.mark.parametrize("mode", list(CompilationMode))
def test_fullgraph_request_checks_compilation_mode(mode):
    result = replace(_resolve(_impl()), compilation_mode=mode)
    context = ExecutionContext(platform="cuda", require_fullgraph=True)
    expected = SupportStatus.UNSUPPORTED if mode is CompilationMode.EAGER_ONLY else SupportStatus.SUPPORTED
    assert result.requested_support(context).status is expected
    assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.SUPPORTED


def test_request_preserves_verified_rejection_reason():
    result = replace(_resolve(_impl()), support=CapabilityResult.unsupported("unsupported head size"))
    assert result.requested_support(ExecutionContext(platform="cuda", require_fullgraph=True)) == result.support


def test_exact_cuda_fa4_dense_path_is_supported():
    context = ExecutionContext(
        platform="cuda",
        require_fullgraph=True,
    )
    result = _resolve(_impl(), context=context)

    assert result.path == "fa4_dense"
    assert result.support.status is SupportStatus.SUPPORTED
    assert result.compilation_mode is CompilationMode.CUSTOM_OP
    assert result.requested_support(context).status is SupportStatus.SUPPORTED


def test_causal_fa4_path_remains_unmigrated():
    result = _resolve(_impl(causal=True))

    assert result.support.status is SupportStatus.UNMIGRATED


def test_fa4_delegates_dimensions_to_kernel(monkeypatch):
    head_dims = (80, 48)
    validator = Mock(return_value=True)
    monkeypatch.setattr(fa, "validate_fa4_head_dims", validator)
    query = torch.empty((1, 16, 8, head_dims[0]), dtype=torch.bfloat16)
    value = torch.empty((1, 16, 8, head_dims[1]), dtype=torch.bfloat16)
    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, query, value, None)
    assert result.support.status is SupportStatus.SUPPORTED
    validator.assert_called_once_with(*head_dims, 8)


def test_fa4_reports_kernel_rejection(monkeypatch):
    monkeypatch.setattr(fa, "validate_fa4_head_dims", Mock(side_effect=AssertionError("kernel dimension constraint")))
    result = _resolve(_impl())
    assert result.support.status is SupportStatus.UNSUPPORTED
    assert "kernel dimension constraint" in result.support.reason
    assert "another backend" in result.support.reason


def test_fa4_missing_validator_remains_unmigrated(monkeypatch):
    monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: False)
    result = _resolve(_impl())
    assert result.support.status is SupportStatus.UNMIGRATED
    assert result.compilation_mode is CompilationMode.EAGER_ONLY


def test_fa4_rejects_mismatched_query_key_dimensions():
    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
    key = torch.empty((1, 16, 8, 32), dtype=torch.bfloat16)
    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, key, key, None)
    assert result.support.status is SupportStatus.UNSUPPORTED
    assert "Q and K head dimensions must match" in result.support.reason


@pytest.mark.parametrize("arch", [90, 100, 120])
def test_fa4_validator_adapter_uses_selected_architecture(monkeypatch, arch):
    validator = Mock()
    monkeypatch.setitem(
        sys.modules,
        "flash_attn.cute.interface",
        SimpleNamespace(_get_device_arch=lambda: arch, _validate_head_dims=validator),
    )
    # Call the real adapter, not the fixture's stand-in.
    verified = _validate_fa4_head_dims(80, 48, 8)
    if arch == 120:
        assert not verified
        validator.assert_not_called()
    else:
        assert verified
        validator.assert_called_once_with(80, 48, arch // 10, 8)


def test_fa4_validator_adapter_handles_missing_private_api(monkeypatch):
    monkeypatch.setitem(sys.modules, "flash_attn.cute.interface", SimpleNamespace())
    assert not _validate_fa4_head_dims(64, 64, 8)


def test_initialized_kernel_identity_overrides_caller_claim():
    result = _resolve(
        _impl(kernel_variant=None),
        context=ExecutionContext(
            platform="cuda",
            kernel_variant="fa4",
        ),
    )

    assert result.support.status is SupportStatus.UNMIGRATED


def test_preconstruction_without_resolved_cuda_kernel_is_unmigrated():
    result = FlashAttentionBackend.resolve_capabilities(ExecutionContext(platform="cuda"))

    assert result.support.status is SupportStatus.UNMIGRATED
    assert result.path == "unverified"


@pytest.mark.parametrize(
    "changes",
    [
        {"platform": "rocm"},
        {"paged_kv": True},
        {"parallel_strategy": ParallelStrategy.ULYSSES},
        {"parallel_strategy": ParallelStrategy.RING},
        {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING},
        {"outer_boundaries": frozenset({OuterBoundary.HSDP})},
    ],
)
def test_unverified_flash_attention_variants_remain_unmigrated(changes):
    result = _resolve(
        _impl(),
        context=replace(ExecutionContext(platform="cuda"), **changes),
    )

    assert result.support.status is SupportStatus.UNMIGRATED
    assert result.compilation_mode is CompilationMode.EAGER_ONLY


def test_unverified_dtype_remains_unmigrated():
    result = _resolve(_impl(), dtype=torch.float32)

    assert result.support.status is SupportStatus.UNMIGRATED


@pytest.mark.parametrize(
    "attn_metadata",
    [
        AttentionMetadata(full_attn_spans=[[(0, 8)]]),
        AttentionMetadata(
            extra={
                "cu_seqlens_q": torch.tensor([0, 16], dtype=torch.int32),
                "cu_seqlens_k": torch.tensor([0, 16], dtype=torch.int32),
                "max_seqlen_q": 16,
                "max_seqlen_k": 16,
            }
        ),
        AttentionMetadata(extra={"kv_cache_dtype": "fp8"}),
    ],
)
def test_metadata_normalization_prevents_dense_support(attn_metadata):
    result = _resolve(_impl(), attn_metadata=attn_metadata)

    assert result.support.status is SupportStatus.UNMIGRATED


def test_unpublished_all_true_mask_remains_runtime_dependent():
    metadata = AttentionMetadata(
        attn_mask=torch.ones((1, 16), dtype=torch.bool),
    )
    result = _resolve(_impl(), attn_metadata=metadata)

    assert result.path == "runtime_mask_dependent"
    assert result.support.status is SupportStatus.UNMIGRATED


def test_producer_published_noop_mask_resolves_dense_fa4():
    metadata = AttentionMetadata(
        attn_mask=torch.ones((1, 16), dtype=torch.bool),
        extra={"attention_mask_mode": "none"},
    )
    result = _resolve(_impl(), attn_metadata=metadata)

    assert result.path == "fa4_dense"
    assert result.support.status is SupportStatus.SUPPORTED


def test_fixed_shape_mask_semantic_change_invalidates_dense_resolution():
    metadata = AttentionMetadata(
        attn_mask=torch.ones((1, 16), dtype=torch.bool),
        extra={"attention_mask_mode": "none"},
    )
    impl = _impl()
    dense = _resolve(impl, attn_metadata=metadata)
    assert dense.support.status is SupportStatus.SUPPORTED

    metadata.attn_mask[:, 8:] = False
    metadata.extra["attention_mask_mode"] = "padding"
    masked = _resolve(impl, attn_metadata=metadata)
    assert masked.support.status is SupportStatus.UNMIGRATED
    request = ExecutionContext(platform="cuda", require_fullgraph=True)
    assert masked.requested_support(request).status is SupportStatus.UNSUPPORTED


def test_piecewise_dispatch_ignores_unused_incomplete_packed_metadata(monkeypatch):
    from vllm_omni.diffusion.attention.backends import flash_attn
    from vllm_omni.diffusion.attention.backends.utils import fa

    metadata = AttentionMetadata(full_attn_spans=[[(0, 16)]], extra={"max_seqlen_q": 16})
    impl = _impl()
    impl.fa_deterministic = False
    impl.softmax_scale = 0.125
    assert _resolve(impl, attn_metadata=metadata).support.status is SupportStatus.UNMIGRATED

    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
    monkeypatch.setattr(fa, "flash_attn_func", lambda *args, **kwargs: query)
    monkeypatch.setattr(flash_attn, "piecewise_attn", lambda *args, **kwargs: query)
    assert impl.forward_cuda(query, query, query, metadata) is query


@pytest.mark.parametrize("input_index", [1, 2])
@pytest.mark.parametrize("mismatch", ["dtype", "device"])
def test_fa4_rejects_incompatible_key_value(input_index, mismatch):
    tensors = [torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) for _ in range(3)]
    tensors[input_index] = (
        tensors[input_index].to(dtype=torch.float16) if mismatch == "dtype" else tensors[input_index].to("meta")
    )
    result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *tensors, None)
    assert result.support.status is SupportStatus.UNSUPPORTED
    assert f"{mismatch}s must match" in result.support.reason


@pytest.mark.parametrize("skip_quant", [False, True])
def test_layer_resolution_uses_effective_kv_quantization(skip_quant):
    from vllm_omni.diffusion.attention.layer import Attention
    from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention

    layer = Attention.__new__(Attention)
    torch.nn.Module.__init__(layer)
    layer.attention = _impl(kernel_variant=None)
    layer._hsdp_compile_boundary_enabled = False
    layer.skip_sequence_parallel = False
    layer._no_parallel_strategy = NoParallelAttention()
    layer.parallel_strategy = layer._no_parallel_strategy
    layer.use_ring = False
    layer.paged_kv_cache_role = None
    layer._kv_cache_dtype = "fp8"
    layer._disable_kv_quant = False
    layer._kv_cache_skip_steps = None
    layer._kv_cache_skip_layers = {0} if skip_quant else None
    layer.layer_idx = 0
    query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16)
    metadata = AttentionMetadata(extra={"kv_cache_dtype": "fp8"})
    result = layer.resolve_execution_path(ExecutionContext(platform="npu"), query, query, query, metadata)
    assert result.path == ("npu_dense" if skip_quant else "npu_unverified")
    assert result.support.status is SupportStatus.UNMIGRATED
    assert metadata.extra == {"kv_cache_dtype": "fp8"}