"""Tests for standardized KernelRegistry (Issue #24)."""

from __future__ import annotations

import pytest

from sagellm_backend.kernels.base import Kernel, KernelRegistry


class DummyKernel(Kernel):
    """Dummy kernel for testing registry behavior."""

    def __init__(self, kernel_name: str, output: str) -> None:
        self._kernel_name = kernel_name
        self._output = output

    @property
    def name(self) -> str:
        return self._kernel_name

    def forward(self, *args: object, **kwargs: object) -> str:  # noqa: ARG002
        return self._output


def test_register_and_query_kernel() -> None:
    """Registry should support operator registration and query."""
    registry = KernelRegistry()
    cpu_linear = DummyKernel("linear_cpu", "cpu_linear")

    registry.register("linear", cpu_linear, provider_type="cpu", priority=10)

    resolved = registry.get("linear", provider_type="cpu")
    assert resolved is cpu_linear


def test_select_highest_priority_within_provider() -> None:
    """Registry should select highest-priority kernel for same provider."""
    registry = KernelRegistry()
    base_cuda = DummyKernel("linear_cuda_base", "cuda_base")
    fast_cuda = DummyKernel("linear_cuda_fast", "cuda_fast")

    registry.register("linear", base_cuda, provider_type="cuda", priority=10)
    registry.register("linear", fast_cuda, provider_type="cuda", priority=100)

    resolved = registry.get("linear", provider_type="cuda")
    assert resolved is fast_cuda


def test_cuda_to_cpu_fallback() -> None:
    """Registry should fallback from CUDA to CPU when CUDA kernel is missing."""
    registry = KernelRegistry()
    cpu_linear = DummyKernel("linear_cpu", "cpu_linear")

    registry.register("linear", cpu_linear, provider_type="cpu", priority=10)

    resolved = registry.get("linear", provider_type="cuda", allow_fallback=True)
    assert resolved is cpu_linear


def test_disable_fallback_raises() -> None:
    """Disabling fallback should fail fast if provider kernel is missing."""
    registry = KernelRegistry()
    cpu_linear = DummyKernel("linear_cpu", "cpu_linear")

    registry.register("linear", cpu_linear, provider_type="cpu", priority=10)

    with pytest.raises(KeyError, match="provider 'cuda'"):
        registry.get("linear", provider_type="cuda", allow_fallback=False)


def test_auto_query_uses_provider_priority_order() -> None:
    """Auto query should follow provider priority order."""
    registry = KernelRegistry(provider_priority=["cuda", "ascend", "cpu"])
    cpu_linear = DummyKernel("linear_cpu", "cpu_linear")
    ascend_linear = DummyKernel("linear_ascend", "ascend_linear")

    registry.register("linear", cpu_linear, provider_type="cpu", priority=1000)
    registry.register("linear", ascend_linear, provider_type="ascend", priority=1)

    resolved = registry.get("linear")
    assert resolved is ascend_linear


def test_list_registered_returns_metadata() -> None:
    """Registry should expose registration metadata for observability/debugging."""
    registry = KernelRegistry()
    cpu_linear = DummyKernel("linear_cpu", "cpu_linear")

    registry.register("linear", cpu_linear, provider_type="cpu", priority=10)

    entries = registry.list_registered("linear")
    assert len(entries) == 1
    assert entries[0].operator_name == "linear"
    assert entries[0].provider_type == "cpu"
    assert entries[0].priority == 10
    assert entries[0].kernel is cpu_linear
