"""Cross-provider capability matrix coverage tests.

Verifies that:
  - All backend providers declare the minimum required capability fields
  - dtype/kernel/graph/stream/collective coverage is complete per provider spec
  - Configurable CPU scenarios (no-int4, no-graph, no-collective) work fail-fast

Issues: #75 #76 #77 #78 #79 #80 #81 #84 #85 #86 #87
"""

from __future__ import annotations

import pytest
from sagellm_protocol import CapabilityDescriptor, DType, KernelKind

from sagellm_backend.providers.cpu import CPUBackendProvider, create_cpu_backend

# ---------------------------------------------------------------------------
# Constants: minimum required capability fields for any provider
# ---------------------------------------------------------------------------

REQUIRED_CAPABILITY_FIELDS = [
    "device_type",
    "supported_dtypes",
    "supported_kernels",
    "has_stream",
    "has_event",
    "has_collective",
    "has_graph_exec",
    "supported_collectives",
    "has_kv_block_copy",
    "has_kv_block_migrate",
]

# Per-provider expected minimum dtypes
CPU_MIN_DTYPES = {DType.FP32, DType.FP16, DType.INT8}

# Per-provider expected minimum kernels
BASE_MIN_KERNELS = {KernelKind.MATMUL, KernelKind.ATTENTION}


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------


def _assert_capability_complete(cap: CapabilityDescriptor, provider_name: str) -> None:
    """Assert that a capability descriptor has all required fields."""
    for field in REQUIRED_CAPABILITY_FIELDS:
        assert hasattr(cap, field), (
            f"[{provider_name}] CapabilityDescriptor missing field '{field}'"
        )
    assert isinstance(cap.supported_dtypes, list), (
        f"[{provider_name}] supported_dtypes must be list"
    )
    assert isinstance(cap.supported_kernels, list), (
        f"[{provider_name}] supported_kernels must be list"
    )
    assert len(cap.supported_dtypes) >= 1, f"[{provider_name}] supported_dtypes must not be empty"
    assert len(cap.supported_kernels) >= 1, f"[{provider_name}] supported_kernels must not be empty"


# ---------------------------------------------------------------------------
# Tests: CPU provider baseline (dtype matrix, #75 #76)
# ---------------------------------------------------------------------------


class TestCPUDtypeMatrix:
    """dtype support matrix for CPU backend (#75 #76)."""

    def test_cpu_declares_fp32_fp16_int8(self) -> None:
        """CPU provider declares FP32/FP16/INT8."""
        backend = CPUBackendProvider()
        cap = backend.capability()

        for dtype in CPU_MIN_DTYPES:
            assert dtype in cap.supported_dtypes, f"CPU missing dtype: {dtype}"

    def test_cpu_rejects_fp8(self) -> None:
        """CPU does not support FP8 and raises NotImplementedError."""
        backend = CPUBackendProvider()
        with pytest.raises(NotImplementedError, match="does not support fp8"):
            backend.capability().require_dtype(DType.FP8, "test")

    def test_cpu_rejects_int4(self) -> None:
        """CPU does not support INT4 and reject kv_alloc fail-fast (#80)."""
        backend = CPUBackendProvider()
        assert DType.INT4 not in backend.capability().supported_dtypes
        with pytest.raises(NotImplementedError, match="does not support int4"):
            backend.kv_block_alloc(num_tokens=32, dtype=DType.INT4)

    @pytest.mark.parametrize("dtype", [DType.FP32, DType.FP16, DType.INT8])
    def test_cpu_accepts_supported_dtypes(self, dtype: DType) -> None:
        """CPU kv_block_alloc succeeds for all declared dtypes."""
        backend = CPUBackendProvider()
        block = backend.kv_block_alloc(num_tokens=16, dtype=dtype)
        assert block.dtype == dtype
        backend.kv_block_free(block)


# ---------------------------------------------------------------------------
# Tests: CPU kernel capability declaration (#77)
# ---------------------------------------------------------------------------


class TestCPUKernelCapability:
    """Kernel capability declaration for CPU backend (#77)."""

    def test_cpu_declares_base_kernels(self) -> None:
        """CPU provider declares MATMUL, ATTENTION, FUSION, LAYERNORM, SOFTMAX."""
        backend = CPUBackendProvider()
        cap = backend.capability()

        for kernel in BASE_MIN_KERNELS:
            assert cap.supports_kernel(kernel), f"CPU missing kernel: {kernel}"

        assert cap.supports_kernel(KernelKind.FUSION)
        assert cap.supports_kernel(KernelKind.LAYERNORM)
        assert cap.supports_kernel(KernelKind.SOFTMAX)

    def test_cpu_capability_complete(self) -> None:
        """CPU CapabilityDescriptor has all required fields."""
        backend = CPUBackendProvider()
        _assert_capability_complete(backend.capability(), "cpu")


# ---------------------------------------------------------------------------
# Tests: graph/stream/collective coverage (#78)
# ---------------------------------------------------------------------------


class TestCPUGraphStreamCollective:
    """graph/stream/collective coverage for CPU backend (#78)."""

    def test_cpu_has_stream_and_event(self) -> None:
        """CPU declares has_stream=True, has_event=True."""
        cap = CPUBackendProvider().capability()
        assert cap.has_stream is True
        assert cap.has_event is True

    def test_cpu_no_graph_exec(self) -> None:
        """CPU does not support graph execution by default."""
        cap = CPUBackendProvider().capability()
        assert cap.has_graph_exec is False
        with pytest.raises(NotImplementedError, match="does not support graph execution"):
            cap.require_graph_exec()

    def test_cpu_no_collective(self) -> None:
        """CPU has no collective operations (single device)."""
        cap = CPUBackendProvider().capability()
        assert cap.has_collective is False
        assert cap.supported_collectives == []
        with pytest.raises(NotImplementedError, match="does not support collective operation"):
            cap.require_collective("all_reduce")

    def test_cpu_stream_create_and_use(self) -> None:
        """CPU stream creation works and synchronize is no-op."""
        backend = CPUBackendProvider()
        stream = backend.create_stream()
        event = backend.create_event()
        stream.synchronize()
        stream.record_event(event)
        event.synchronize()


# ---------------------------------------------------------------------------
# Tests: Configurable capability matrix (#79 #84)
# ---------------------------------------------------------------------------


class TestConfigurableCapabilityMatrix:
    """Configurable capability matrix for CPU backend (#79 #84)."""

    def test_override_with_dict_single_field(self) -> None:
        """Override a single field via dict."""
        backend = CPUBackendProvider(capabilities={"has_graph_exec": True})
        assert backend.capability().has_graph_exec is True

    def test_override_dtypes_list(self) -> None:
        """Override supported_dtypes via dict."""
        backend = CPUBackendProvider(capabilities={"supported_dtypes": [DType.FP32]})
        cap = backend.capability()
        assert cap.supported_dtypes == [DType.FP32]
        assert DType.FP16 not in cap.supported_dtypes

    def test_override_with_full_descriptor(self) -> None:
        """Override using a full CapabilityDescriptor object."""
        custom = CapabilityDescriptor(
            device_name="custom-test-device",
            device_type="cpu-custom",
            supported_dtypes=[DType.FP32],
            supported_kernels=[KernelKind.MATMUL],
            has_stream=True,
            has_event=True,
            has_collective=True,
            has_graph_exec=True,
            supported_collectives=["all_reduce"],
            has_kv_block_copy=True,
            has_kv_block_migrate=True,
        )
        backend = CPUBackendProvider(capabilities=custom)
        cap = backend.capability()
        assert cap.device_type == "cpu-custom"
        assert cap.has_collective is True
        assert "all_reduce" in cap.supported_collectives

    def test_create_cpu_backend_with_config(self) -> None:
        """create_cpu_backend() factory propagates config dict."""
        backend = create_cpu_backend(config={"supported_dtypes": [DType.FP32, DType.FP16]})
        cap = backend.capability()
        assert DType.FP32 in cap.supported_dtypes
        assert DType.FP16 in cap.supported_dtypes
        assert DType.INT8 not in cap.supported_dtypes


# ---------------------------------------------------------------------------
# Tests: No-int4 CI scenario (#80 #85 #87)
# ---------------------------------------------------------------------------


class TestNoInt4CIScenario:
    """No-int4 CI scenario tests (#80 #85 #87)."""

    def test_cpu_default_no_int4(self) -> None:
        """Default CPU backend does not support INT4."""
        backend = CPUBackendProvider()
        assert DType.INT4 not in backend.capability().supported_dtypes

    def test_configured_no_int4_kv_alloc_raises(self) -> None:
        """Explicit no-INT4 configuration causes kv_block_alloc to raise."""
        backend = CPUBackendProvider(capabilities={"supported_dtypes": [DType.FP32, DType.FP16]})
        with pytest.raises(NotImplementedError, match="does not support int4"):
            backend.kv_block_alloc(num_tokens=32, dtype=DType.INT4)

    def test_configured_no_int4_supports_dtype_check(self) -> None:
        """supports_dtype returns False for INT4 in no-int4 config."""
        backend = CPUBackendProvider(capabilities={"supported_dtypes": [DType.FP32]})
        assert backend.capability().supports_dtype(DType.INT4) is False

    def test_require_dtype_raises_for_int4(self) -> None:
        """require_dtype raises NotImplementedError for INT4."""
        backend = CPUBackendProvider()
        with pytest.raises(NotImplementedError, match="does not support int4"):
            backend.capability().require_dtype(DType.INT4, "test_op")


# ---------------------------------------------------------------------------
# Tests: No-graph CI scenario (#81 #86)
# ---------------------------------------------------------------------------


class TestNoGraphCIScenario:
    """No-graph CI scenario tests (#81 #86)."""

    def test_cpu_default_no_graph(self) -> None:
        """Default CPU backend does not support graph execution."""
        backend = CPUBackendProvider()
        assert backend.capability().has_graph_exec is False

    def test_configured_no_graph_raises(self) -> None:
        """Explicit has_graph_exec=False configuration raises on require_graph_exec."""
        backend = CPUBackendProvider(capabilities={"has_graph_exec": False})
        with pytest.raises(NotImplementedError, match="does not support graph execution"):
            backend.capability().require_graph_exec()

    def test_configured_with_graph_does_not_raise(self) -> None:
        """Explicit has_graph_exec=True configuration does not raise on require_graph_exec."""
        backend = CPUBackendProvider(capabilities={"has_graph_exec": True})
        # Should not raise
        backend.capability().require_graph_exec()


# ---------------------------------------------------------------------------
# Tests: CapabilityDescriptor helper methods completeness
# ---------------------------------------------------------------------------


class TestCapabilityDescriptorHelpers:
    """Tests for CapabilityDescriptor helper methods."""

    def test_supports_dtype_true_for_declared(self) -> None:
        backend = CPUBackendProvider()
        assert backend.capability().supports_dtype(DType.FP32) is True

    def test_supports_dtype_false_for_undeclared(self) -> None:
        backend = CPUBackendProvider()
        assert backend.capability().supports_dtype(DType.INT4) is False
        assert backend.capability().supports_dtype(DType.FP8) is False

    def test_supports_kernel_true_for_declared(self) -> None:
        backend = CPUBackendProvider()
        assert backend.capability().supports_kernel(KernelKind.MATMUL) is True
        assert backend.capability().supports_kernel(KernelKind.ATTENTION) is True

    def test_require_dtype_no_raise_for_supported(self) -> None:
        backend = CPUBackendProvider()
        # Should not raise
        backend.capability().require_dtype(DType.FP32, "matmul")
        backend.capability().require_dtype(DType.FP16, "attention")

    def test_require_collective_raises_for_cpu(self) -> None:
        backend = CPUBackendProvider()
        with pytest.raises(NotImplementedError, match="does not support collective operation"):
            backend.capability().require_collective("all_reduce")
