"""Regression tests for A05 stability fixes.

Covers issues:
- #55: ROCmAttentionBackend must fail-fast when paged attention is requested
- #59: @triton.jit at module import time must not raise OSError (lazy init)
- #88: ProfileCache.clear() must work correctly
- #89: CudaInspector must have single consistent implementation
"""

from __future__ import annotations

import tempfile
from pathlib import Path
from unittest import mock

import pytest
import torch

from sagellm_backend.attention.base import AttentionMetadata

# ---------------------------------------------------------------------------
# Issue #59 — Triton @jit lazy-init: import must not crash
# ---------------------------------------------------------------------------


class TestTritonLazyInit:
    """Regression: importing fused_ops must not raise OSError (Issue #59)."""

    def test_fused_ops_import_succeeds(self):
        """Importing fused_ops module must never raise OSError at import time."""
        # This exercises the fix: @triton.jit is now wrapped in @functools.cache
        # factory functions so it runs lazily, not at module import time.
        try:
            from sagellm_backend.kernels import fused_ops  # noqa: F401
        except OSError as exc:
            pytest.fail(
                f"[#59 regression] fused_ops import raised OSError — "
                f"@triton.jit must not be applied at module import time: {exc}"
            )

    def test_fused_silu_mul_kernel_import_succeeds(self):
        """FusedSiLUMulKernel must be importable without crashing."""
        try:
            from sagellm_backend.kernels import FusedSiLUMulKernel  # noqa: F401
        except OSError as exc:
            pytest.fail(f"[#59 regression] FusedSiLUMulKernel import failed: {exc}")

    def test_fused_rmsnorm_kernel_import_succeeds(self):
        """FusedRMSNormKernel must be importable without crashing."""
        try:
            from sagellm_backend.kernels import FusedRMSNormKernel  # noqa: F401
        except OSError as exc:
            pytest.fail(f"[#59 regression] FusedRMSNormKernel import failed: {exc}")

    def test_sagellm_backend_top_level_import_succeeds(self):
        """Top-level sagellm_backend import must not crash."""
        try:
            import sagellm_backend  # noqa: F401
        except OSError as exc:
            pytest.fail(f"[#59 regression] sagellm_backend import raised OSError: {exc}")

    def test_fused_silu_mul_cpu_works_after_import(self):
        """FusedSiLUMulKernel must produce correct output on CPU after lazy init."""
        from sagellm_backend.kernels import FusedSiLUMulKernel

        kernel = FusedSiLUMulKernel()
        gate = torch.randn(4, 64)
        up = torch.randn(4, 64)

        out = kernel(gate, up)

        expected = torch.nn.functional.silu(gate) * up
        assert torch.allclose(out, expected, rtol=1e-5, atol=1e-6), (
            "[#59 regression] FusedSiLUMulKernel CPU output mismatch after lazy init"
        )


# ---------------------------------------------------------------------------
# Issue #55 — ROCmAttentionBackend must fail-fast for paged attention
# ---------------------------------------------------------------------------


def _make_mock_device_props():
    """Return a mock device properties object for tests."""
    props = mock.Mock()
    props.name = "Hygon DCU K100"
    props.total_memory = 64 * 1024 * 1024 * 1024  # 64 GB
    props.major = 9
    props.minor = 0
    return props


@pytest.fixture()
def rocm_backend():
    """Build a ROCmAttentionBackend with all CUDA calls mocked out."""
    with (
        mock.patch("torch.cuda.is_available", return_value=True),
        mock.patch("torch.cuda.get_device_properties", return_value=_make_mock_device_props()),
    ):
        from sagellm_backend.attention.rocm import ROCmAttentionBackend

        return ROCmAttentionBackend(device="cuda:0", use_flash_attn=False)


class TestROCmPagedAttentionFailFast:
    """Regression: ROCm backend must declare and enforce paged attention limits (#55)."""

    def test_supports_paged_attention_is_false(self, rocm_backend):
        """ROCmAttentionBackend.supports_paged_attention must be False."""
        assert rocm_backend.supports_paged_attention is False, (
            "[#55 regression] ROCmAttentionBackend.supports_paged_attention must be False "
            "until ROCm PagedAttention is implemented."
        )

    def test_capabilities_reflects_paged_attention_false(self, rocm_backend):
        """get_capabilities() must expose supports_paged_attention=False."""
        with (
            mock.patch("torch.cuda.is_available", return_value=True),
            mock.patch(
                "torch.cuda.get_device_properties",
                return_value=_make_mock_device_props(),
            ),
        ):
            caps = rocm_backend.get_capabilities()
        assert caps["supports_paged_attention"] is False, (
            "[#55 regression] get_capabilities() must report supports_paged_attention=False"
        )

    def test_forward_with_block_tables_raises_not_implemented(self, rocm_backend):
        """forward() must raise NotImplementedError when block_tables != None (fail-fast)."""

        batch, seq_len, num_heads, head_dim = 1, 4, 2, 8
        query = torch.randn(batch, seq_len, num_heads, head_dim)
        key = torch.randn(batch, seq_len, num_heads, head_dim)
        value = torch.randn(batch, seq_len, num_heads, head_dim)

        # Simulate a paged-attention call by providing block_tables
        block_tables = torch.zeros(batch, 1, dtype=torch.long)
        metadata = AttentionMetadata(
            seq_lens=[seq_len] * batch,
            block_tables=block_tables,
            num_prefill_tokens=batch * seq_len,
            num_decode_tokens=0,
        )

        with pytest.raises(NotImplementedError, match="PagedAttention"):
            rocm_backend.forward(query, key, value, metadata)

    def test_forward_without_block_tables_does_not_raise(self, rocm_backend):
        """forward() with block_tables=None must not raise (standard attention path)."""

        batch, seq_len, num_heads, head_dim = 1, 4, 2, 8
        query = torch.randn(batch, seq_len, num_heads, head_dim)
        key = torch.randn(batch, seq_len, num_heads, head_dim)
        value = torch.randn(batch, seq_len, num_heads, head_dim)

        metadata = AttentionMetadata(
            seq_lens=[seq_len] * batch,
            block_tables=None,
            num_prefill_tokens=batch * seq_len,
            num_decode_tokens=0,
        )

        # With mocked cuda device, forward will try to move tensors to "cuda:0"
        # and then call pytorch SDPA — skip execution errors, only assert no
        # NotImplementedError about paged attention is raised.
        try:
            rocm_backend.forward(query, key, value, metadata)
        except NotImplementedError as exc:
            if "PagedAttention" in str(exc):
                pytest.fail(
                    "[#55 regression] forward() raised NotImplementedError about "
                    f"PagedAttention even though block_tables=None: {exc}"
                )
            # Other NotImplementedError (e.g. device ops) are acceptable here
        except Exception:
            # CUDA device errors are expected in mock environment — not our concern
            pass


# ---------------------------------------------------------------------------
# Issue #88 — ProfileCache.clear() regression
# ---------------------------------------------------------------------------


class TestProfileCacheClearRegression:
    """Regression: ProfileCache must have a working clear() method (Issue #88)."""

    def test_clear_deletes_existing_cache_file(self):
        """clear() must delete the cache file if it exists."""
        from sagellm_backend.hardware.cache import ProfileCache

        with tempfile.TemporaryDirectory() as tmpdir:
            cache_path = Path(tmpdir) / "hardware_profile.json"
            cache_path.write_text('{"version":"1.0","profiles":{}}')

            cache = ProfileCache(cache_path=cache_path)
            assert cache_path.exists()

            cache.clear()
            assert not cache_path.exists(), (
                "[#88 regression] ProfileCache.clear() did not delete the cache file"
            )

    def test_clear_on_nonexistent_file_is_noop(self):
        """clear() on a non-existent cache must not raise any exception."""
        from sagellm_backend.hardware.cache import ProfileCache

        with tempfile.TemporaryDirectory() as tmpdir:
            cache_path = Path(tmpdir) / "nonexistent.json"
            cache = ProfileCache(cache_path=cache_path)

            try:
                cache.clear()
            except Exception as exc:
                pytest.fail(
                    f"[#88 regression] ProfileCache.clear() raised on non-existent file: {exc}"
                )

    def test_clear_method_exists(self):
        """ProfileCache must have a 'clear' attribute that is callable."""
        from sagellm_backend.hardware.cache import ProfileCache

        assert hasattr(ProfileCache, "clear"), "[#88 regression] ProfileCache has no 'clear' method"
        assert callable(ProfileCache.clear), "[#88 regression] ProfileCache.clear is not callable"


# ---------------------------------------------------------------------------
# Issue #89 — CudaInspector single consistent implementation
# ---------------------------------------------------------------------------


class TestCudaInspectorRegression:
    """Regression: CudaInspector must have one consistent implementation (Issue #89)."""

    def test_single_cuda_inspector_class_in_module(self):
        """There must be exactly one CudaInspector class in the cuda module."""
        import inspect

        from sagellm_backend.hardware.inspectors import cuda as cuda_mod

        cuda_inspector_classes = [
            name
            for name, obj in inspect.getmembers(cuda_mod, inspect.isclass)
            if name == "CudaInspector"
        ]
        assert len(cuda_inspector_classes) == 1, (
            f"[#89 regression] Expected 1 CudaInspector class, "
            f"found {len(cuda_inspector_classes)}. Duplicate definition causes runtime mismatch."
        )

    @mock.patch("subprocess.run")
    def test_l1_5_returns_operational_message(self, mock_run):
        """inspect_l1_5_management_tools must return 'nvidia-smi operational' message."""
        from sagellm_backend.hardware.inspectors.cuda import CudaInspector

        mock_result = mock.Mock()
        mock_result.returncode = 0
        mock_result.stdout = "535.0\n"

        mock_gpu_result = mock.Mock()
        mock_gpu_result.returncode = 0
        mock_gpu_result.stdout = "NVIDIA A100\nNVIDIA A100\n"

        mock_run.side_effect = [mock_result, mock_gpu_result]

        inspector = CudaInspector()
        result = inspector.inspect_l1_5_management_tools()

        assert result.available, "[#89 regression] Expected result.available=True"
        assert "nvidia-smi operational" in result.message, (
            f"[#89 regression] Expected 'nvidia-smi operational' in message, got: {result.message!r}"
        )
        assert "driver_version" in result.details, (
            f"[#89 regression] Expected 'driver_version' in details, got: {result.details}"
        )
        assert result.details["driver_version"] == "535.0", (
            f"[#89 regression] Expected driver_version='535.0', got: {result.details['driver_version']}"
        )
        assert result.details["status"] == "operational", (
            f"[#89 regression] Expected status='operational', got: {result.details['status']}"
        )
