"""Tests for kernel injection framework.

This module tests the kernel injection infrastructure implemented
for issue #36.
"""

from __future__ import annotations

import pytest

from sagellm_core.kernel_injection import (
    BackendKernelRegistry,
    KernelInjector,
    KernelSpec,
    OperatorRegistry,
    register_backend_kernels,
)
from sagellm_protocol import OperatorConstraint, OperatorType


# Test fixtures
def dummy_linear_kernel(*args, **kwargs):
    """Dummy linear kernel for testing."""
    return "linear_result"


def dummy_attention_kernel(*args, **kwargs):
    """Dummy attention kernel for testing."""
    return "attention_result"


def dummy_rmsnorm_kernel(*args, **kwargs):
    """Dummy RMSNorm kernel for testing."""
    return "rmsnorm_result"


class DummyLayer:
    """Dummy layer for testing kernel injection."""

    def __init__(self):
        self._backend = None
        self._kernel = None

    def set_backend(self, backend):
        """Set backend."""
        self._backend = backend

    def set_kernel(self, kernel):
        """Set kernel."""
        self._kernel = kernel


# Tests for KernelSpec
def test_kernel_spec_creation():
    """Test KernelSpec creation."""
    constraint = OperatorConstraint(device_type="cuda", dtype="fp16")
    spec = KernelSpec(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint,
        priority=10,
        description="Test kernel",
    )

    assert spec.operator_type == OperatorType.LINEAR
    assert spec.kernel_impl == dummy_linear_kernel
    assert spec.constraint == constraint
    assert spec.priority == 10
    assert spec.description == "Test kernel"


def test_kernel_spec_matches():
    """Test KernelSpec matching logic."""
    constraint = OperatorConstraint(
        device_type="cuda",
        dtype="fp16",
        min_compute_capability="8.0",
    )
    spec = KernelSpec(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint,
    )

    # Should match
    assert spec.matches(device_type="cuda", dtype="fp16", compute_capability="8.0")
    assert spec.matches(device_type="cuda", dtype="fp16", compute_capability="8.6")

    # Should not match
    assert not spec.matches(device_type="cpu", dtype="fp16")
    assert not spec.matches(device_type="cuda", dtype="fp32")
    assert not spec.matches(device_type="cuda", dtype="fp16", compute_capability="7.0")


def test_kernel_spec_matches_partial():
    """Test KernelSpec matching with partial constraints."""
    # Only device type constraint
    constraint = OperatorConstraint(device_type="cuda")
    spec = KernelSpec(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint,
    )

    # Should match any dtype on CUDA
    assert spec.matches(device_type="cuda", dtype="fp16")
    assert spec.matches(device_type="cuda", dtype="fp32")
    assert not spec.matches(device_type="cpu", dtype="fp16")


# Tests for OperatorRegistry
def test_operator_registry_creation():
    """Test OperatorRegistry creation."""
    registry = OperatorRegistry()
    assert len(registry.list_operators()) == 0
    assert registry.get_kernel_count() == 0


def test_operator_registry_register_kernel():
    """Test kernel registration."""
    registry = OperatorRegistry()

    constraint = OperatorConstraint(device_type="cuda", dtype="fp16")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint,
        priority=10,
    )

    assert len(registry.list_operators()) == 1
    assert OperatorType.LINEAR in registry.list_operators()
    assert registry.get_kernel_count() == 1
    assert registry.get_kernel_count(OperatorType.LINEAR) == 1


def test_operator_registry_query_kernel():
    """Test kernel querying."""
    registry = OperatorRegistry()

    constraint = OperatorConstraint(device_type="cuda", dtype="fp16")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint,
        priority=10,
    )

    # Query should find matching kernel
    spec = registry.query_kernel(OperatorType.LINEAR, device_type="cuda", dtype="fp16")
    assert spec is not None
    assert spec.kernel_impl == dummy_linear_kernel

    # Query should not find non-matching kernel
    spec = registry.query_kernel(OperatorType.LINEAR, device_type="cpu", dtype="fp16")
    assert spec is None


def test_operator_registry_priority():
    """Test kernel priority ordering."""
    registry = OperatorRegistry()

    constraint1 = OperatorConstraint(device_type="cuda", dtype="fp16")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=lambda: "low_priority",
        constraint=constraint1,
        priority=5,
    )

    constraint2 = OperatorConstraint(device_type="cuda", dtype="fp16")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=lambda: "high_priority",
        constraint=constraint2,
        priority=10,
    )

    # Should return the higher priority kernel
    spec = registry.query_kernel(OperatorType.LINEAR, device_type="cuda", dtype="fp16")
    assert spec is not None
    assert spec.priority == 10
    assert spec.kernel_impl() == "high_priority"


def test_operator_registry_query_all():
    """Test querying all matching kernels."""
    registry = OperatorRegistry()

    for i in range(3):
        constraint = OperatorConstraint(device_type="cuda", dtype="fp16")
        registry.register_kernel(
            operator_type=OperatorType.LINEAR,
            kernel_impl=lambda: f"kernel_{i}",
            constraint=constraint,
            priority=i,
        )

    specs = registry.query_all_kernels(OperatorType.LINEAR, device_type="cuda", dtype="fp16")
    assert len(specs) == 3
    # Should be ordered by priority (descending)
    assert specs[0].priority == 2
    assert specs[1].priority == 1
    assert specs[2].priority == 0


# Tests for KernelInjector
def test_kernel_injector_creation():
    """Test KernelInjector creation."""
    registry = OperatorRegistry()
    injector = KernelInjector(registry)

    assert injector._registry == registry
    assert injector._backend is None


def test_kernel_injector_inject_into_layer():
    """Test injecting kernel into layer."""
    registry = OperatorRegistry()

    constraint = OperatorConstraint(device_type="cuda", dtype="fp16")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint,
    )

    injector = KernelInjector(registry)
    layer = DummyLayer()

    # Should successfully inject
    success = injector.inject_into_layer(
        layer, OperatorType.LINEAR, device_type="cuda", dtype="fp16"
    )

    assert success is True
    assert layer._kernel == dummy_linear_kernel


def test_kernel_injector_no_matching_kernel():
    """Test injection when no matching kernel found."""
    registry = OperatorRegistry()
    injector = KernelInjector(registry)
    layer = DummyLayer()

    # Should fail gracefully
    success = injector.inject_into_layer(
        layer, OperatorType.LINEAR, device_type="cuda", dtype="fp16"
    )

    assert success is False
    assert layer._kernel is None


def test_kernel_injector_fallback_kernel():
    """Test fallback kernel injection when no backend kernel matches."""
    registry = OperatorRegistry()
    injector = KernelInjector(registry)
    layer = DummyLayer()

    success = injector.inject_into_layer(
        layer,
        OperatorType.LINEAR,
        fallback_kernel=dummy_linear_kernel,
        device_type="cuda",
        dtype="fp16",
    )

    assert success is False
    assert layer._kernel == dummy_linear_kernel


def test_kernel_injector_inject_into_model():
    """Test injecting kernels into model."""

    class DummyModel:
        def __init__(self):
            self.linear = DummyLayer()
            self.attention = DummyLayer()

        def named_modules(self):
            return [
                ("linear", self.linear),
                ("attention", self.attention),
            ]

    registry = OperatorRegistry()

    # Register linear kernel
    constraint_linear = OperatorConstraint(device_type="cuda")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint_linear,
    )

    # Register attention kernel
    constraint_attention = OperatorConstraint(device_type="cuda")
    registry.register_kernel(
        operator_type=OperatorType.ATTENTION,
        kernel_impl=dummy_attention_kernel,
        constraint=constraint_attention,
    )

    injector = KernelInjector(registry)
    model = DummyModel()

    operator_type_map = {
        "linear": OperatorType.LINEAR,
        "attention": OperatorType.ATTENTION,
    }

    results = injector.inject_into_model(model, operator_type_map, device_type="cuda")

    assert len(results) == 2
    assert results["linear"] is True
    assert results["attention"] is True
    assert model.linear._kernel == dummy_linear_kernel
    assert model.attention._kernel == dummy_attention_kernel


# Tests for BackendKernelRegistry
def test_backend_kernel_registry():
    """Test BackendKernelRegistry helper."""

    class DummyBackend:
        pass

    registry = OperatorRegistry()
    backend = DummyBackend()
    bridge = BackendKernelRegistry(registry, backend)

    # Register linear kernel
    bridge.register_linear_kernel(
        dummy_linear_kernel,
        device_type="cuda",
        dtype="fp16",
        priority=10,
    )

    # Should be registered in registry
    spec = registry.query_kernel(OperatorType.LINEAR, device_type="cuda", dtype="fp16")
    assert spec is not None
    assert spec.kernel_impl == dummy_linear_kernel
    assert spec.priority == 10


def test_backend_kernel_registry_multiple_kernels():
    """Test registering multiple kernel types."""

    class DummyBackend:
        pass

    registry = OperatorRegistry()
    backend = DummyBackend()
    bridge = BackendKernelRegistry(registry, backend)

    # Register multiple kernels
    bridge.register_linear_kernel(dummy_linear_kernel, device_type="cuda")
    bridge.register_attention_kernel(dummy_attention_kernel, device_type="cuda")
    bridge.register_rmsnorm_kernel(dummy_rmsnorm_kernel, device_type="cuda")

    # All should be registered
    assert registry.get_kernel_count() == 3
    assert len(registry.list_operators()) == 3


def test_operator_registry_replace_kernel():
    """Test runtime operator kernel replacement."""
    registry = OperatorRegistry()
    constraint = OperatorConstraint(device_type="cuda", dtype="fp16")

    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=lambda *_args, **_kwargs: "old",
        constraint=constraint,
        priority=5,
    )

    previous = registry.replace_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=lambda *_args, **_kwargs: "new",
        constraint=constraint,
        priority=20,
    )

    assert previous is not None
    current = registry.query_kernel(OperatorType.LINEAR, device_type="cuda", dtype="fp16")
    assert current is not None
    assert current.priority == 20
    assert current.kernel_impl() == "new"


def test_register_backend_kernels_bridge():
    """Test bridging backend.get_kernel() into OperatorRegistry."""

    class DummyBackend:
        def capability(self):
            class Capability:
                device_type = "cpu"

            return Capability()

        def get_kernel(self, name: str):
            kernels = {
                "linear": dummy_linear_kernel,
                "embedding": lambda *args, **kwargs: "embedding_result",
            }
            if name not in kernels:
                raise KeyError(name)
            return kernels[name]

    registry = OperatorRegistry()
    count = register_backend_kernels(DummyBackend(), registry)

    assert count >= 2
    assert registry.query_kernel(OperatorType.LINEAR, device_type="cpu") is not None
    assert registry.query_kernel(OperatorType.EMBEDDING, device_type="cpu") is not None


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