"""Tests for flash attention configuration and attention backend telemetry (EDG-48)."""

from __future__ import annotations

import time
from dataclasses import dataclass
from unittest.mock import MagicMock, patch


from octomil.serve import (
    EchoBackend,
    InferenceBackend,
    InferenceMetrics,
    LlamaCppBackend,
    MLXBackend,
)
from octomil.telemetry import TelemetryReporter


# ---------------------------------------------------------------------------
# InferenceMetrics — attention_backend field
# ---------------------------------------------------------------------------


class TestInferenceMetricsAttentionBackend:
    def test_default_is_standard(self):
        m = InferenceMetrics()
        assert m.attention_backend == "standard"

    def test_custom_attention_backend(self):
        m = InferenceMetrics(attention_backend="flash_attention")
        assert m.attention_backend == "flash_attention"

    def test_metal_fused(self):
        m = InferenceMetrics(attention_backend="metal_fused")
        assert m.attention_backend == "metal_fused"

    def test_sdpa(self):
        m = InferenceMetrics(attention_backend="sdpa")
        assert m.attention_backend == "sdpa"


# ---------------------------------------------------------------------------
# InferenceBackend — attention_backend class attribute
# ---------------------------------------------------------------------------


class TestInferenceBackendAttentionBackend:
    def test_base_backend_default(self):
        assert InferenceBackend.attention_backend == "standard"

    def test_mlx_backend_metal_fused(self):
        assert MLXBackend.attention_backend == "metal_fused"

    def test_llamacpp_backend_flash_attention(self):
        assert LlamaCppBackend.attention_backend == "flash_attention"

    def test_echo_backend_inherits_standard(self):
        echo = EchoBackend()
        assert echo.attention_backend == "standard"


# ---------------------------------------------------------------------------
# LlamaCppBackend — flash_attn=True in Llama() constructor
# ---------------------------------------------------------------------------


class TestLlamaCppFlashAttn:
    def test_load_model_local_gguf_passes_flash_attn(self):
        """Llama() constructor should receive flash_attn=True for local .gguf files."""
        mock_llama_cls = MagicMock()
        mock_llama_cls.return_value = MagicMock()

        mock_llama_module = MagicMock(Llama=mock_llama_cls, LlamaCache=MagicMock())
        with patch.dict("sys.modules", {"llama_cpp": mock_llama_module}):
            backend = LlamaCppBackend(cache_enabled=False)
            backend.load_model("model.gguf")

        mock_llama_cls.assert_called_once()
        call_kwargs = mock_llama_cls.call_args
        assert (
            call_kwargs.kwargs.get("flash_attn") is True
        ), "Llama() constructor must receive flash_attn=True for local .gguf files"
        assert call_kwargs.kwargs.get("model_path") == "model.gguf"

    def test_load_model_hf_repo_passes_flash_attn(self):
        """Llama.from_pretrained() should receive flash_attn=True for HF repo IDs."""
        mock_llama_cls = MagicMock()
        mock_from_pretrained = MagicMock(return_value=MagicMock())
        mock_llama_cls.from_pretrained = mock_from_pretrained

        mock_llama_module = MagicMock(Llama=mock_llama_cls, LlamaCache=MagicMock())
        with patch.dict("sys.modules", {"llama_cpp": mock_llama_module}):
            backend = LlamaCppBackend(cache_enabled=False)
            backend.load_model("user/some-model")

        mock_from_pretrained.assert_called_once()
        call_kwargs = mock_from_pretrained.call_args
        assert (
            call_kwargs.kwargs.get("flash_attn") is True
        ), "Llama.from_pretrained() must receive flash_attn=True for HF repo IDs"
        assert call_kwargs.kwargs.get("repo_id") == "user/some-model"

    def test_load_model_resolved_gguf_passes_flash_attn(self):
        """Llama.from_pretrained() should receive flash_attn=True for resolver-resolved models."""
        mock_llama_cls = MagicMock()
        mock_from_pretrained = MagicMock(return_value=MagicMock())
        mock_llama_cls.from_pretrained = mock_from_pretrained

        mock_resolved = MagicMock()
        mock_resolved.is_gguf = True
        mock_resolved.filename = "model-Q4_K_M.gguf"
        mock_resolved.hf_repo = "test-org/test-model"

        mock_llama_module = MagicMock(Llama=mock_llama_cls, LlamaCache=MagicMock())
        with (
            patch.dict("sys.modules", {"llama_cpp": mock_llama_module}),
            patch("octomil.serve._resolve_new", return_value=mock_resolved),
        ):
            backend = LlamaCppBackend(cache_enabled=False)
            backend.load_model("short-name")

        mock_from_pretrained.assert_called_once()
        call_kwargs = mock_from_pretrained.call_args
        assert (
            call_kwargs.kwargs.get("flash_attn") is True
        ), "Llama.from_pretrained() must receive flash_attn=True for resolver-resolved models"

    def test_load_model_legacy_catalog_passes_flash_attn(self):
        """Llama.from_pretrained() should receive flash_attn=True for legacy catalog models."""
        mock_llama_cls = MagicMock()
        mock_from_pretrained = MagicMock(return_value=MagicMock())
        mock_llama_cls.from_pretrained = mock_from_pretrained

        mock_llama_module = MagicMock(Llama=mock_llama_cls, LlamaCache=MagicMock())
        with (
            patch.dict("sys.modules", {"llama_cpp": mock_llama_module}),
            patch(
                "octomil.serve._resolve_new",
                side_effect=__import__(
                    "octomil.models.resolver", fromlist=["ModelResolutionError"]
                ).ModelResolutionError("not found"),
            ),
            patch(
                "octomil.serve._GGUF_MODELS", {"test-model": ("org/repo", "file.gguf")}
            ),
        ):
            backend = LlamaCppBackend(cache_enabled=False)
            backend.load_model("test-model")

        mock_from_pretrained.assert_called_once()
        call_kwargs = mock_from_pretrained.call_args
        assert (
            call_kwargs.kwargs.get("flash_attn") is True
        ), "Llama.from_pretrained() must receive flash_attn=True for legacy catalog models"


# ---------------------------------------------------------------------------
# MLXBackend — Metal fused attention comment and attribute
# ---------------------------------------------------------------------------


class TestMLXAttentionBackend:
    def test_mlx_engine_module_docstring_mentions_metal_fused(self):
        """The mlx_engine module docstring should document Metal fused attention."""
        from octomil.engines import mlx_engine

        assert "Metal fused attention" in mlx_engine.__doc__

    def test_mlx_backend_class_docstring_mentions_metal_fused(self):
        assert "metal_fused" in MLXBackend.__doc__


# ---------------------------------------------------------------------------
# ORT backend — attention_backend attribute
# ---------------------------------------------------------------------------


class TestORTAttentionBackend:
    def test_ort_backend_attention_backend_is_sdpa(self):
        from octomil.engines.ort_engine import _ORTBackend

        backend = _ORTBackend("test-model")
        assert backend.attention_backend == "sdpa"

    def test_ort_generate_session_returns_sdpa(self):
        """ORT session generate should set attention_backend='sdpa' in metrics."""
        import numpy as np

        from octomil.engines.ort_engine import _ORTBackend

        mock_input = MagicMock()
        mock_input.name = "input_ids"
        mock_input.shape = [1, 10]
        mock_input.type = "tensor(float)"

        mock_session = MagicMock()
        mock_session.get_inputs.return_value = [mock_input]
        mock_session.run.return_value = [np.array([[0.5]])]

        backend = _ORTBackend("test-model")
        backend._session = mock_session
        backend._use_genai = False

        @dataclass
        class FakeRequest:
            messages: list[dict[str, str]]
            max_tokens: int = 32

        request = FakeRequest(messages=[{"role": "user", "content": "Hi"}])
        _text, metrics = backend.generate(request)
        assert metrics.attention_backend == "sdpa"

    def test_ort_generate_genai_returns_sdpa(self):
        """ORT GenAI generate should set attention_backend='sdpa' in metrics."""
        from octomil.engines.ort_engine import _ORTBackend

        mock_og = MagicMock()
        mock_model = MagicMock()
        mock_tokenizer = MagicMock()
        mock_tokenizer.encode.return_value = [1, 2, 3]
        mock_tokenizer.decode.return_value = "Hello"

        mock_generator = MagicMock()
        mock_generator.is_done.side_effect = [False, False, True]
        mock_generator.get_next_tokens.return_value = [42]
        mock_og.Generator.return_value = mock_generator
        mock_og.GeneratorParams.return_value = MagicMock()

        backend = _ORTBackend("test-model")
        backend._model = mock_model
        backend._tokenizer = mock_tokenizer
        backend._use_genai = True

        @dataclass
        class FakeRequest:
            messages: list[dict[str, str]]
            max_tokens: int = 512

        request = FakeRequest(messages=[{"role": "user", "content": "Hi"}])

        with patch.dict("sys.modules", {"onnxruntime_genai": mock_og}):
            _text, metrics = backend.generate(request)

        assert metrics.attention_backend == "sdpa"


# ---------------------------------------------------------------------------
# TelemetryReporter — attention_backend in payloads
# ---------------------------------------------------------------------------


class TestTelemetryAttentionBackend:
    def test_inference_started_includes_attention_backend(self):
        sent: list[dict] = []

        def mock_send(client, url, headers, payload):
            sent.append(payload)

        with patch.object(TelemetryReporter, "_send", side_effect=mock_send):
            reporter = TelemetryReporter(
                api_key="key",
                api_base="https://api.test.com/api/v1",
                org_id="org-1",
                device_id="dev-1",
            )
            reporter.report_inference_started(
                model_id="model-a",
                version="1.0",
                session_id="s1",
                attention_backend="flash_attention",
            )
            time.sleep(0.15)
            reporter.close()

        assert len(sent) >= 1
        event = sent[0]["events"][0]
        assert event["name"] == "inference.started"
        assert event["attributes"]["inference.attention_backend"] == "flash_attention"

    def test_inference_started_no_attention_backend_has_no_attr(self):
        sent: list[dict] = []

        def mock_send(client, url, headers, payload):
            sent.append(payload)

        with patch.object(TelemetryReporter, "_send", side_effect=mock_send):
            reporter = TelemetryReporter(
                api_key="key",
                api_base="https://api.test.com/api/v1",
                org_id="org-1",
                device_id="dev-1",
            )
            reporter.report_inference_started(
                model_id="model-a",
                version="1.0",
                session_id="s1",
            )
            time.sleep(0.15)
            reporter.close()

        assert len(sent) >= 1
        attrs = sent[0]["events"][0]["attributes"]
        assert "inference.attention_backend" not in attrs

    def test_inference_completed_includes_attention_backend(self):
        sent: list[dict] = []

        def mock_send(client, url, headers, payload):
            sent.append(payload)

        with patch.object(TelemetryReporter, "_send", side_effect=mock_send):
            reporter = TelemetryReporter(
                api_key="key",
                api_base="https://api.test.com/api/v1",
                org_id="org-1",
                device_id="dev-1",
            )
            reporter.report_inference_completed(
                session_id="s1",
                model_id="model-a",
                version="1.0",
                total_chunks=10,
                total_duration_ms=500.0,
                ttfc_ms=30.0,
                throughput=20.0,
                attention_backend="metal_fused",
            )
            time.sleep(0.15)
            reporter.close()

        assert len(sent) >= 1
        event = sent[0]["events"][0]
        assert event["name"] == "inference.completed"
        attrs = event["attributes"]
        assert attrs["inference.attention_backend"] == "metal_fused"
        # Original attributes should still be present
        assert attrs["inference.total_tokens"] == 10
        assert attrs["inference.throughput_tps"] == 20.0

    def test_inference_completed_without_attention_backend(self):
        """When attention_backend is not passed, it should not appear in attributes."""
        sent: list[dict] = []

        def mock_send(client, url, headers, payload):
            sent.append(payload)

        with patch.object(TelemetryReporter, "_send", side_effect=mock_send):
            reporter = TelemetryReporter(
                api_key="key",
                api_base="https://api.test.com/api/v1",
                org_id="org-1",
                device_id="dev-1",
            )
            reporter.report_inference_completed(
                session_id="s1",
                model_id="model-a",
                version="1.0",
                total_chunks=5,
                total_duration_ms=250.0,
                ttfc_ms=15.0,
                throughput=10.0,
            )
            time.sleep(0.15)
            reporter.close()

        assert len(sent) >= 1
        attrs = sent[0]["events"][0]["attributes"]
        assert "inference.attention_backend" not in attrs


# ---------------------------------------------------------------------------
# Engine plugins — attention_backend awareness
# ---------------------------------------------------------------------------


class TestEnginePluginAttentionBackend:
    def test_llamacpp_engine_creates_backend_with_flash_attention(self):
        """LlamaCppEngine.create_backend should return a backend with flash_attention."""
        from octomil.engines.llamacpp_engine import LlamaCppEngine

        engine = LlamaCppEngine()
        with patch("octomil.serve.LlamaCppBackend") as MockBackend:
            mock_instance = MagicMock()
            mock_instance.attention_backend = "flash_attention"
            MockBackend.return_value = mock_instance
            backend = engine.create_backend("test-model")
            assert backend.attention_backend == "flash_attention"

    def test_mlx_engine_creates_backend_with_metal_fused(self):
        """MLXEngine.create_backend should return a backend with metal_fused."""
        from octomil.engines.mlx_engine import MLXEngine

        engine = MLXEngine()
        with patch("octomil.serve.MLXBackend") as MockBackend:
            mock_instance = MagicMock()
            mock_instance.attention_backend = "metal_fused"
            MockBackend.return_value = mock_instance
            backend = engine.create_backend("test-model")
            assert backend.attention_backend == "metal_fused"

    def test_ort_engine_creates_backend_with_sdpa(self):
        """ONNXRuntimeEngine.create_backend should return a backend with sdpa."""
        from octomil.engines.ort_engine import ONNXRuntimeEngine

        engine = ONNXRuntimeEngine()
        backend = engine.create_backend("test-model")
        assert backend.attention_backend == "sdpa"
