"""Tests for feature variant management (Issue #32)."""

from __future__ import annotations

import torch

from sagellm_backend.attention.base import AttentionMetadata
from sagellm_backend.feature_variants import FeatureVariantManager


def test_attention_variant_decode_kv_prefers_paged() -> None:
    manager = FeatureVariantManager.with_default_rules()
    metadata = AttentionMetadata.for_decode(
        context_lens=[16, 32],
        block_tables=torch.zeros((2, 2), dtype=torch.int32),
        device="cpu",
    )

    decision = manager.select_attention(metadata, {"cpu", "flash", "paged"})
    assert decision.variant_name == "decode_kv_cache"
    assert decision.implementation_name == "paged"


def test_attention_variant_prefill_no_kv_prefers_flash_then_cpu() -> None:
    manager = FeatureVariantManager.with_default_rules()
    metadata = AttentionMetadata.for_prefill(
        seq_lens=[64, 128],
        device="cpu",
        extra={"use_kv_cache": False},
    )

    flash_decision = manager.select_attention(metadata, {"cpu", "flash"})
    assert flash_decision.variant_name == "prefill_no_cache"
    assert flash_decision.implementation_name == "flash"

    cpu_decision = manager.select_attention(metadata, {"cpu"})
    assert cpu_decision.implementation_name == "cpu"


def test_activation_variant_moe_masked_prefers_fused() -> None:
    manager = FeatureVariantManager.with_default_rules()

    decision = manager.select(
        category="activation",
        context={"moe_masked": True},
        available_impls={"silu", "fused_silu_mul"},
    )

    assert decision.variant_name == "activation_moe_masked"
    assert decision.implementation_name == "fused_silu_mul"


def test_gemm_variant_block_wise_quant_prefers_quant_kernel() -> None:
    manager = FeatureVariantManager.with_default_rules()

    decision = manager.select(
        category="gemm",
        context={"quantization_mode": "block_wise"},
        available_impls={"linear", "linear_block_wise_quant"},
    )

    assert decision.variant_name == "gemm_block_wise_quant"
    assert decision.implementation_name == "linear_block_wise_quant"


def test_manual_override_requires_availability() -> None:
    manager = FeatureVariantManager.with_default_rules()
    metadata = AttentionMetadata.for_prefill(seq_lens=[32], device="cpu")

    decision = manager.select_attention(metadata, {"cpu", "flash"}, preferred_backend="cpu")
    assert decision.variant_name == "manual_override"
    assert decision.implementation_name == "cpu"
