# SPDX-License-Identifier: Apache-2.0
"""Unit tests for accuracy benchmark orchestration."""

import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from omlx.admin.accuracy_benchmark import (
    VALID_BENCHMARKS,
    AccuracyBenchmarkRequest,
    AccuracyBenchmarkRun,
    _accumulated_results,
    add_to_queue,
    cleanup_old_runs,
    create_run,
    get_accumulated_results,
    get_queue_status,
    get_run,
    reset_accumulated_results,
    run_accuracy_benchmark,
)
from omlx.model_settings import ModelSettings


class TestAccuracyBenchmarkRequest:
    def test_valid_request(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 300, "gsm8k": 100},
        )
        assert req.model_id == "test-model"
        assert "mmlu" in req.benchmarks
        assert req.benchmarks["gsm8k"] == 100

    def test_full_dataset_size_zero(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 0},
        )
        assert req.benchmarks["mmlu"] == 0

    def test_empty_benchmarks_rejected(self):
        with pytest.raises(Exception):
            AccuracyBenchmarkRequest(
                model_id="test-model",
                benchmarks={},
            )

    def test_invalid_benchmark_rejected(self):
        with pytest.raises(Exception):
            AccuracyBenchmarkRequest(
                model_id="test-model",
                benchmarks={"invalid_bench": 100},
            )

    def test_all_valid_benchmarks(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={b: 100 for b in VALID_BENCHMARKS},
        )
        assert len(req.benchmarks) == len(VALID_BENCHMARKS)

    def test_enable_thinking_default_false(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        assert req.enable_thinking is False

    def test_enable_thinking_true(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
            enable_thinking=True,
        )
        assert req.enable_thinking is True

    def test_sampling_profile_default_deterministic(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        assert req.sampling_profile == "deterministic"

    def test_sampling_profile_model_settings_accepted(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
            sampling_profile="model_settings",
        )
        assert req.sampling_profile == "model_settings"

    def test_sampling_profile_invalid_rejected(self):
        with pytest.raises(Exception):
            AccuracyBenchmarkRequest(
                model_id="test-model",
                benchmarks={"mmlu": 100},
                sampling_profile="wild",
            )


class TestQueueAndResults:
    def setup_method(self):
        from omlx.admin.accuracy_benchmark import _queue
        _queue.clear()
        reset_accumulated_results()

    def test_add_to_queue(self):
        req = AccuracyBenchmarkRequest(
            model_id="model-a",
            benchmarks={"mmlu": 100},
        )
        add_to_queue(req)
        status = get_queue_status()
        assert len(status["queue"]) == 1
        assert status["queue"][0]["model_id"] == "model-a"

    def test_queue_status_empty(self):
        status = get_queue_status()
        assert status["running"] is False
        assert len(status["queue"]) == 0

    def test_accumulated_results(self):
        _accumulated_results.append({"model_id": "m1", "benchmark": "mmlu", "accuracy": 0.5})
        results = get_accumulated_results()
        assert len(results) == 1
        assert results[0]["model_id"] == "m1"

    def test_reset_accumulated_results(self):
        _accumulated_results.append({"model_id": "m1", "benchmark": "mmlu", "accuracy": 0.5})
        reset_accumulated_results()
        assert len(get_accumulated_results()) == 0


class TestRunLifecycle:
    def setup_method(self):
        from omlx.admin.accuracy_benchmark import _accuracy_runs
        _accuracy_runs.clear()

    def test_create_run(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        run = create_run(req)
        assert run.bench_id is not None
        assert run.status == "running"
        assert run.request == req

    def test_get_run(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        run = create_run(req)
        found = get_run(run.bench_id)
        assert found is run

    def test_get_run_not_found(self):
        assert get_run("nonexistent") is None

    def test_cleanup_old_runs(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        run1 = create_run(req)
        run2 = create_run(req)
        run1.status = "completed"
        run2.status = "running"

        cleanup_old_runs()

        assert get_run(run1.bench_id) is None
        assert get_run(run2.bench_id) is run2

    def test_cleanup_error_runs(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        run = create_run(req)
        run.status = "error"

        cleanup_old_runs()
        assert get_run(run.bench_id) is None


class TestRunAccuracyBenchmark:
    @pytest.mark.asyncio
    async def test_sends_done_event(self):
        """Verify that a successful run sends a done event."""
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        run = create_run(req)

        # Mock engine_pool
        mock_engine = AsyncMock()
        mock_engine.chat = AsyncMock(return_value=MagicMock(text="A"))

        mock_pool = MagicMock()
        mock_pool.get_loaded_model_ids = MagicMock(return_value=[])
        mock_pool.get_engine = AsyncMock(return_value=mock_engine)
        mock_pool._unload_engine = AsyncMock()

        # Mock evaluator
        mock_result = MagicMock()
        mock_result.benchmark_name = "mmlu"
        mock_result.accuracy = 0.75
        mock_result.total_questions = 4
        mock_result.correct_count = 3
        mock_result.time_seconds = 1.0
        mock_result.category_scores = None
        mock_result.thinking_used = False

        mock_evaluator = MagicMock()
        mock_evaluator.load_dataset = AsyncMock(return_value=[{"id": "1"}])
        mock_evaluator.run = AsyncMock(return_value=mock_result)

        mock_bench_cls = MagicMock(return_value=mock_evaluator)

        with patch.dict("omlx.eval.BENCHMARKS", {"mmlu": mock_bench_cls}, clear=True):
            await run_accuracy_benchmark(run, mock_pool)

        # Collect all events from the replay log.
        events = list(run.events)

        event_types = [e["type"] for e in events]
        assert "done" in event_types
        assert run.status == "completed"

    @pytest.mark.asyncio
    async def test_cancellation(self):
        """Verify that cancelling stops the run."""
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 100},
        )
        run = create_run(req)
        run.status = "cancelled"  # Pre-cancel

        mock_pool = MagicMock()
        mock_pool.get_loaded_model_ids = MagicMock(return_value=[])
        mock_pool.get_engine = AsyncMock(return_value=MagicMock())
        mock_pool._unload_engine = AsyncMock()

        mock_evaluator = MagicMock()
        mock_evaluator.load_dataset = AsyncMock(return_value=[])
        mock_evaluator.run = AsyncMock(return_value=MagicMock(
            benchmark_name="mmlu",
            accuracy=0.0,
            total_questions=0,
            correct_count=0,
            time_seconds=0.0,
            category_scores=None,
        ))

        mock_bench_cls = MagicMock(return_value=mock_evaluator)

        with patch.dict("omlx.eval.BENCHMARKS", {"mmlu": mock_bench_cls}):
            await run_accuracy_benchmark(run, mock_pool)

        # Should have stopped early
        assert len(run.results) == 0


class TestSamplingProfile:
    """sampling_profile gates whether per-model sampling reaches the evaluator.

    Default "deterministic" must read nothing (reproducible greedy scores);
    "model_settings" must forward the model's configured sampling. See #606 /
    the #1254 deterministic-default request.
    """

    def _mock_pool(self, model_settings):
        mock_engine = AsyncMock()
        mock_engine.chat = AsyncMock(return_value=MagicMock(text="A"))
        mock_pool = MagicMock()
        mock_pool.get_loaded_model_ids = MagicMock(return_value=[])
        mock_pool.get_engine = AsyncMock(return_value=mock_engine)
        mock_pool._unload_engine = AsyncMock()
        mock_pool._settings_manager.get_settings = MagicMock(return_value=model_settings)
        return mock_pool

    async def _captured_sampling_kwargs(self, req, mock_pool):
        run = create_run(req)
        mock_result = MagicMock(
            benchmark_name="mmlu", accuracy=0.5, total_questions=1,
            correct_count=1, time_seconds=0.1, category_scores=None,
            thinking_used=False,
        )
        mock_evaluator = MagicMock()
        mock_evaluator.load_dataset = AsyncMock(return_value=[{"id": "1"}])
        mock_evaluator.run = AsyncMock(return_value=mock_result)
        mock_bench_cls = MagicMock(return_value=mock_evaluator)
        with patch.dict("omlx.eval.BENCHMARKS", {"mmlu": mock_bench_cls}, clear=True):
            await run_accuracy_benchmark(run, mock_pool)
        return mock_evaluator.run.call_args.kwargs["sampling_kwargs"]

    @pytest.mark.asyncio
    async def test_deterministic_ignores_model_settings(self):
        # Default profile is "deterministic".
        req = AccuracyBenchmarkRequest(model_id="test-model", benchmarks={"mmlu": 1})
        mock_pool = self._mock_pool(ModelSettings(temperature=0.9, top_p=0.95))
        assert await self._captured_sampling_kwargs(req, mock_pool) == {}

    @pytest.mark.asyncio
    async def test_model_settings_forwards_sampling(self):
        req = AccuracyBenchmarkRequest(
            model_id="test-model",
            benchmarks={"mmlu": 1},
            sampling_profile="model_settings",
        )
        mock_pool = self._mock_pool(ModelSettings(temperature=0.9, top_p=0.95))
        sampling_kwargs = await self._captured_sampling_kwargs(req, mock_pool)
        assert sampling_kwargs["temperature"] == 0.9
        assert sampling_kwargs["top_p"] == 0.95

    @pytest.mark.asyncio
    async def test_deterministic_keeps_chat_template_kwargs(self):
        # Template kwargs are prompt construction, not sampling — forwarded
        # even under the deterministic profile.
        req = AccuracyBenchmarkRequest(model_id="test-model", benchmarks={"mmlu": 1})
        mock_pool = self._mock_pool(
            ModelSettings(temperature=0.9, chat_template_kwargs={"custom_flag": True})
        )
        sampling_kwargs = await self._captured_sampling_kwargs(req, mock_pool)
        assert sampling_kwargs == {"chat_template_kwargs": {"custom_flag": True}}


# =============================================================================
# External endpoint accuracy benchmark tests
# =============================================================================


def _external_dict():
    return {
        "base_url": "http://localhost:8001/v1",
        "api_key": "sk-test",
        "model": "remote-model",
    }


class TestExternalAccuracyRequest:
    def test_external_accepted(self):
        req = AccuracyBenchmarkRequest(
            model_id="remote-model",
            benchmarks={"mmlu": 100},
            external=_external_dict(),
        )
        assert req.external is not None
        assert req.external.model == "remote-model"

    def test_external_forces_thinking_off(self):
        req = AccuracyBenchmarkRequest(
            model_id="remote-model",
            benchmarks={"mmlu": 100},
            enable_thinking=True,
            external=_external_dict(),
        )
        assert req.enable_thinking is False

    def test_local_keeps_thinking(self):
        req = AccuracyBenchmarkRequest(
            model_id="local-model",
            benchmarks={"mmlu": 100},
            enable_thinking=True,
        )
        assert req.enable_thinking is True

    def test_queue_status_flags_external(self):
        req = AccuracyBenchmarkRequest(
            model_id="remote-model",
            benchmarks={"mmlu": 100},
            external=_external_dict(),
        )
        add_to_queue(req)
        try:
            entry = get_queue_status()["queue"][-1]
            assert entry["external"] is True
        finally:
            from omlx.admin.accuracy_benchmark import _queue

            _queue.clear()


class TestExternalAccuracyRun:
    def _mock_result(self):
        return MagicMock(
            benchmark_name="mmlu",
            accuracy=0.5,
            total_questions=2,
            correct_count=1,
            time_seconds=0.1,
            category_scores=None,
            thinking_used=False,
            question_results=[],
        )

    def _mock_evaluator(self):
        mock_evaluator = MagicMock()
        mock_evaluator.load_dataset = AsyncMock(return_value=[{"id": "1"}])
        mock_evaluator.run = AsyncMock(return_value=self._mock_result())
        return mock_evaluator

    def _external_request(self):
        return AccuracyBenchmarkRequest(
            model_id="remote-model",
            benchmarks={"mmlu": 100},
            batch_size=4,
            external=_external_dict(),
        )

    @pytest.mark.asyncio
    async def test_external_run_uses_adapter_and_skips_pool(self):
        run = create_run(self._external_request())
        mock_pool = MagicMock()
        mock_evaluator = self._mock_evaluator()
        mock_bench_cls = MagicMock(return_value=mock_evaluator)

        mock_adapter = MagicMock()
        mock_adapter.preflight = AsyncMock()
        mock_client = MagicMock()
        mock_client.aclose = AsyncMock()

        with (
            patch.dict("omlx.eval.BENCHMARKS", {"mmlu": mock_bench_cls}, clear=True),
            patch(
                "omlx.admin.accuracy_benchmark.ExternalAPIClient",
                return_value=mock_client,
            ),
            patch(
                "omlx.admin.accuracy_benchmark.ExternalChatAdapter",
                return_value=mock_adapter,
            ) as adapter_cls,
        ):
            await run_accuracy_benchmark(run, mock_pool)

        assert run.status == "completed"
        mock_pool.get_engine.assert_not_called()
        mock_pool.get_loaded_model_ids.assert_not_called()
        mock_pool._unload_engine.assert_not_called()
        mock_adapter.preflight.assert_awaited_once()
        adapter_cls.assert_called_once_with(mock_client, "deterministic")
        # Evaluator got the adapter, empty sampling kwargs, thinking off
        call = mock_evaluator.run.call_args
        assert call.args[0] is mock_adapter
        assert call.kwargs["sampling_kwargs"] == {}
        assert call.kwargs["enable_thinking"] is False
        assert call.kwargs["batch_size"] == 4
        # Result carries the external flag and the remote model name
        assert run.results[0]["external"] is True
        assert run.results[0]["model_id"] == "remote-model"
        mock_client.aclose.assert_awaited()
        # Clean up accumulated results this test appended
        reset_accumulated_results()

    @pytest.mark.asyncio
    async def test_external_preflight_failure_emits_error(self):
        from omlx.admin.external_api import ExternalEndpointError

        run = create_run(self._external_request())
        mock_pool = MagicMock()

        mock_adapter = MagicMock()
        mock_adapter.preflight = AsyncMock(
            side_effect=ExternalEndpointError(
                "External endpoint rejected the API key (HTTP 401)"
            )
        )
        mock_client = MagicMock()
        mock_client.aclose = AsyncMock()

        with (
            patch(
                "omlx.admin.accuracy_benchmark.ExternalAPIClient",
                return_value=mock_client,
            ),
            patch(
                "omlx.admin.accuracy_benchmark.ExternalChatAdapter",
                return_value=mock_adapter,
            ),
        ):
            await run_accuracy_benchmark(run, mock_pool)

        assert run.status == "error"
        assert "rejected the API key" in run.error_message
        error_events = [e for e in run.events if e["type"] == "error"]
        assert error_events
        mock_client.aclose.assert_awaited()