"""Tests for fused kernels integration.

This module tests the fused kernel integration implemented for issue #39.
"""

from __future__ import annotations

import pytest
import torch

from sagellm_core.kernel_injection import (
    FusedKernelManager,
    FusionConfig,
    OperatorRegistry,
    detect_fusion_opportunities,
)
from sagellm_protocol import OperatorType


# Test fixtures
class MockBackend:
    """Mock backend for testing."""

    def __init__(self, has_fused_kernels: bool = True):
        self._has_fused_kernels = has_fused_kernels
        self._kernels = {}

        if has_fused_kernels:
            # Register mock fused kernels
            self._kernels["fused_silu_mul"] = self._mock_fused_silu_mul
            self._kernels["fused_rmsnorm"] = self._mock_fused_rmsnorm
            self._kernels["fused_qkv_projection"] = self._mock_fused_qkv_projection

    def capability(self):
        """Return mock capability."""

        class MockCapability:
            device_name = "CUDA"

        return MockCapability()

    def get_kernel(self, name: str):
        """Get kernel by name."""
        if name not in self._kernels:
            raise KeyError(f"Kernel '{name}' not found")
        return self._kernels[name]

    @staticmethod
    def _mock_fused_silu_mul(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
        """Mock fused SiLU * Mul kernel."""
        return torch.nn.functional.silu(gate) * up

    @staticmethod
    def _mock_fused_rmsnorm(
        x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
    ) -> torch.Tensor:
        """Mock fused RMSNorm kernel."""
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + eps)
        return x * weight

    @staticmethod
    def _mock_fused_qkv_projection(
        hidden_states: torch.Tensor,
        qkv_weight: torch.Tensor,
        num_heads: int,
        num_kv_heads: int,
        head_dim: int,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Mock fused QKV projection kernel."""
        # Simplified mock implementation
        batch_size, seq_len, hidden_size = hidden_states.shape
        q = torch.randn(batch_size, seq_len, num_heads * head_dim)
        k = torch.randn(batch_size, seq_len, num_kv_heads * head_dim)
        v = torch.randn(batch_size, seq_len, num_kv_heads * head_dim)
        return q, k, v


class MockModel:
    """Mock model for testing fusion detection."""

    def __init__(self):
        self.mlp = self.MLP()
        self.attention = self.Attention()
        self.norm = self.RMSNorm()

    def named_modules(self):
        """Return named modules."""
        yield "mlp", self.mlp
        yield "attention", self.attention
        yield "norm", self.norm

    class MLP:
        """Mock MLP with gate and up projections."""

        def __init__(self):
            self.gate_proj = torch.nn.Linear(512, 2048)
            self.up_proj = torch.nn.Linear(512, 2048)
            self.down_proj = torch.nn.Linear(2048, 512)

    class Attention:
        """Mock attention with Q/K/V projections."""

        def __init__(self):
            self.q_proj = torch.nn.Linear(512, 512)
            self.k_proj = torch.nn.Linear(512, 512)
            self.v_proj = torch.nn.Linear(512, 512)
            self.o_proj = torch.nn.Linear(512, 512)

    class RMSNorm:
        """Mock RMSNorm."""

        def __init__(self):
            self.weight = torch.nn.Parameter(torch.ones(512))


# Tests for FusionConfig
def test_fusion_config_defaults():
    """Test FusionConfig default values."""
    config = FusionConfig()
    assert config.enabled is True
    assert config.enable_silu_mul is True
    assert config.enable_rmsnorm is True
    assert config.enable_qkv_projection is True
    assert config.auto_detect is True
    assert config.priority_boost == 50


def test_fusion_config_custom():
    """Test FusionConfig with custom values."""
    config = FusionConfig(
        enabled=False,
        enable_silu_mul=False,
        priority_boost=100,
    )
    assert config.enabled is False
    assert config.enable_silu_mul is False
    assert config.priority_boost == 100


# Tests for FusedKernelManager
def test_fused_kernel_manager_creation():
    """Test FusedKernelManager creation."""
    backend = MockBackend()
    config = FusionConfig()
    manager = FusedKernelManager(backend, config)

    assert manager._backend == backend
    assert manager._config == config
    assert len(manager._patterns) == 3  # silu_mul, rmsnorm, qkv_projection


def test_fusion_patterns_initialization():
    """Test fusion patterns are initialized correctly."""
    backend = MockBackend()
    manager = FusedKernelManager(backend)

    patterns = manager._patterns
    assert "silu_mul" in patterns
    assert "rmsnorm" in patterns
    assert "qkv_projection" in patterns

    assert patterns["silu_mul"].fused_operator_type == OperatorType.FUSED_SILU_MUL
    assert patterns["rmsnorm"].fused_operator_type == OperatorType.FUSED_RMSNORM
    assert patterns["qkv_projection"].fused_operator_type == OperatorType.FUSED_QKV_PROJECTION


def test_setup_fusion_success():
    """Test successful fusion setup."""
    backend = MockBackend(has_fused_kernels=True)
    manager = FusedKernelManager(backend)
    registry = OperatorRegistry()

    results = manager.setup_fusion(registry)

    # All three fused kernels should be registered
    assert results["silu_mul"] is True
    assert results["rmsnorm"] is True
    assert results["qkv_projection"] is True

    # Check registry has the kernels
    assert registry.get_kernel_count() == 3


def test_setup_fusion_disabled():
    """Test fusion setup when disabled."""
    backend = MockBackend()
    config = FusionConfig(enabled=False)
    manager = FusedKernelManager(backend, config)
    registry = OperatorRegistry()

    results = manager.setup_fusion(registry)

    assert len(results) == 0
    assert registry.get_kernel_count() == 0


def test_setup_fusion_selective_enable():
    """Test selective fusion enable."""
    backend = MockBackend()
    config = FusionConfig(
        enabled=True,
        enable_silu_mul=True,
        enable_rmsnorm=False,
        enable_qkv_projection=False,
    )
    manager = FusedKernelManager(backend, config)
    registry = OperatorRegistry()

    results = manager.setup_fusion(registry)

    # Only silu_mul should be registered
    assert results.get("silu_mul") is True
    assert results.get("rmsnorm") is False
    assert results.get("qkv_projection") is False


def test_can_fuse():
    """Test can_fuse method."""
    backend = MockBackend()
    manager = FusedKernelManager(backend)
    registry = OperatorRegistry()

    # Before setup, fusion not available
    assert manager.can_fuse("silu_mul") is False

    # After setup, fusion available
    manager.setup_fusion(registry)
    assert manager.can_fuse("silu_mul") is True
    assert manager.can_fuse("rmsnorm") is True
    assert manager.can_fuse("qkv_projection") is True


def test_can_fuse_disabled():
    """Test can_fuse when fusion is disabled."""
    backend = MockBackend()
    config = FusionConfig(enabled=False)
    manager = FusedKernelManager(backend, config)
    registry = OperatorRegistry()

    manager.setup_fusion(registry)

    # Fusion disabled, should return False
    assert manager.can_fuse("silu_mul") is False


def test_get_fusion_stats():
    """Test get_fusion_stats method."""
    backend = MockBackend()
    manager = FusedKernelManager(backend)
    registry = OperatorRegistry()

    manager.setup_fusion(registry)
    stats = manager.get_fusion_stats()

    assert stats["enabled"] is True
    assert stats["total_patterns"] == 3
    assert stats["registered_count"] == 3
    assert "silu_mul" in stats["patterns"]
    assert stats["patterns"]["silu_mul"]["registered"] is True


def test_list_available_fusions():
    """Test list_available_fusions method."""
    backend = MockBackend()
    manager = FusedKernelManager(backend)
    registry = OperatorRegistry()

    manager.setup_fusion(registry)
    fusions = manager.list_available_fusions()

    assert len(fusions) == 3
    assert "silu_mul" in fusions
    assert "rmsnorm" in fusions
    assert "qkv_projection" in fusions


def test_list_available_fusions_selective():
    """Test list_available_fusions with selective enable."""
    backend = MockBackend()
    config = FusionConfig(
        enabled=True,
        enable_silu_mul=True,
        enable_rmsnorm=False,
        enable_qkv_projection=False,
    )
    manager = FusedKernelManager(backend, config)
    registry = OperatorRegistry()

    manager.setup_fusion(registry)
    fusions = manager.list_available_fusions()

    assert len(fusions) == 1
    assert "silu_mul" in fusions
    assert "rmsnorm" not in fusions


# Tests for detect_fusion_opportunities
def test_detect_fusion_opportunities():
    """Test fusion opportunity detection."""
    model = MockModel()
    opportunities = detect_fusion_opportunities(model)

    # Should detect opportunities
    assert "silu_mul" in opportunities
    assert "rmsnorm" in opportunities
    assert "qkv_projection" in opportunities

    # Should find MLP for silu_mul
    assert len(opportunities["silu_mul"]) > 0
    assert "mlp" in opportunities["silu_mul"][0]

    # Should find attention for qkv_projection
    assert len(opportunities["qkv_projection"]) > 0
    assert "attention" in opportunities["qkv_projection"][0]

    # Should find norm for rmsnorm
    assert len(opportunities["rmsnorm"]) > 0
    assert "norm" in opportunities["rmsnorm"][0]


# Integration tests
def test_end_to_end_integration():
    """Test end-to-end fused kernel integration."""
    # Setup
    backend = MockBackend()
    config = FusionConfig(enabled=True)
    manager = FusedKernelManager(backend, config)
    registry = OperatorRegistry()

    # Register fused kernels
    results = manager.setup_fusion(registry)
    assert sum(results.values()) == 3

    # Check availability
    assert manager.can_fuse("silu_mul", device="cuda")

    # Get stats
    stats = manager.get_fusion_stats()
    assert stats["registered_count"] == 3

    # List available
    fusions = manager.list_available_fusions()
    assert len(fusions) == 3


def test_backend_without_fused_kernels():
    """Test behavior when backend doesn't have fused kernels."""

    class BackendWithoutFusion:
        def capability(self):
            class MockCapability:
                device_name = "CPU"

            return MockCapability()

        # No get_kernel method

    backend = BackendWithoutFusion()
    manager = FusedKernelManager(backend)
    registry = OperatorRegistry()

    results = manager.setup_fusion(registry)

    # Should handle gracefully
    assert len(results) == 0
    assert registry.get_kernel_count() == 0


# Performance comparison tests (mock)
def test_fusion_performance_tracking():
    """Test that fusion config allows performance tracking."""
    config = FusionConfig(
        min_speedup_ratio=1.2,  # Require 20% speedup
        priority_boost=100,
    )

    assert config.min_speedup_ratio == 1.2
    assert config.priority_boost == 100


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
