"""Ticket #20: memory-aware model ladder, conditional backend order, retry.

All behavior is exercised with mocked probes and loaders: no GPU, no
network, no hardware. The threshold tests pin the per-backend rung thresholds
to the ticket #19 installed-build measurement (issue #19, comment 12913): the
transcribe-time VRAM delta of each rung plus the 512 MiB protocol headroom.
Re-measuring changes those constants, and this file must change with them.
"""

import os
from contextlib import ExitStack
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

from vocalinux.speech_recognition import model_ladder
from vocalinux.speech_recognition.model_ladder import (
    DEFAULT_RUNG,
    FLOOR_RUNG_ENGLISH,
    FLOOR_RUNG_MULTILINGUAL,
    HEADROOM_MIB,
    LOW_MEMORY_RUNG,
    PREFER_CUDA_Q5_0_ABOVE_MIB,
    PREFER_CUDA_Q8_0_ABOVE_MIB,
    PREFER_Q5_0_ABOVE_MIB,
    PREFER_Q8_0_ABOVE_MIB,
    cuda_toolkit_ready_from_probe,
    next_rung_down,
    preferred_backend_order,
    resolve_startup_rung,
    select_ladder_rung,
)
from vocalinux.speech_recognition.recognition_manager import SpeechRecognitionManager
from vocalinux.utils.whispercpp_model_info import WHISPERCPP_MODEL_INFO


def _size_mb(name: str) -> float:
    return float(WHISPERCPP_MODEL_INFO[name]["size_mb"])


class TestThresholdDerivation:
    """Thresholds are the ticket #19 recorded transcribe-time deltas plus headroom.

    The inputs are copied from issue #19, comment 12913 ("Per-rung load
    footprint", transcribe delta column), not recomputed from the code.
    """

    #: Transcribe-time VRAM delta (MiB), peak minus before, per backend and rung.
    VULKAN_Q8_0_DELTA_MIB: float = 1132.0
    VULKAN_Q5_0_DELTA_MIB: float = 847.0
    CUDA_Q8_0_DELTA_MIB: float = 1178.0
    CUDA_Q5_0_DELTA_MIB: float = 892.0

    def test_headroom_is_the_protocol_value(self) -> None:
        assert HEADROOM_MIB == 512.0

    def test_vulkan_q8_threshold_is_measured_delta_plus_headroom(self) -> None:
        assert PREFER_Q8_0_ABOVE_MIB == self.VULKAN_Q8_0_DELTA_MIB + HEADROOM_MIB
        assert PREFER_Q8_0_ABOVE_MIB == 1644.0

    def test_vulkan_q5_threshold_is_measured_delta_plus_headroom(self) -> None:
        assert PREFER_Q5_0_ABOVE_MIB == self.VULKAN_Q5_0_DELTA_MIB + HEADROOM_MIB
        assert PREFER_Q5_0_ABOVE_MIB == 1359.0

    def test_cuda_q8_threshold_is_measured_delta_plus_headroom(self) -> None:
        assert PREFER_CUDA_Q8_0_ABOVE_MIB == self.CUDA_Q8_0_DELTA_MIB + HEADROOM_MIB
        assert PREFER_CUDA_Q8_0_ABOVE_MIB == 1690.0

    def test_cuda_q5_threshold_is_measured_delta_plus_headroom(self) -> None:
        assert PREFER_CUDA_Q5_0_ABOVE_MIB == self.CUDA_Q5_0_DELTA_MIB + HEADROOM_MIB
        assert PREFER_CUDA_Q5_0_ABOVE_MIB == 1404.0

    def test_rungs_are_ordered_on_each_backend(self) -> None:
        assert PREFER_Q8_0_ABOVE_MIB > PREFER_Q5_0_ABOVE_MIB
        assert PREFER_CUDA_Q8_0_ABOVE_MIB > PREFER_CUDA_Q5_0_ABOVE_MIB
        assert _size_mb(DEFAULT_RUNG) > _size_mb(LOW_MEMORY_RUNG)


class TestSelectLadderRung:
    """Free memory picks the rung; the floor pairs with the language."""

    def test_plenty_of_memory_keeps_default(self) -> None:
        assert select_ladder_rung(True, PREFER_Q8_0_ABOVE_MIB) == DEFAULT_RUNG
        assert select_ladder_rung(False, 99999.0) == DEFAULT_RUNG

    def test_between_thresholds_steps_to_q5(self) -> None:
        assert select_ladder_rung(True, PREFER_Q8_0_ABOVE_MIB - 1) == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, PREFER_Q5_0_ABOVE_MIB) == LOW_MEMORY_RUNG

    def test_below_q5_fits_the_floor_by_language(self) -> None:
        assert select_ladder_rung(True, PREFER_Q5_0_ABOVE_MIB - 1) == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(False, PREFER_Q5_0_ABOVE_MIB - 1) == FLOOR_RUNG_MULTILINGUAL
        assert select_ladder_rung(True, 0.0) == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(False, 0.0) == FLOOR_RUNG_MULTILINGUAL

    def test_recorded_boundaries_select_recorded_rungs(self) -> None:
        assert select_ladder_rung(True, 1644.0) == DEFAULT_RUNG
        assert select_ladder_rung(True, 1500.0) == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, 1359.0) == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, 1358.0) == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(True, 1200.0) == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(False, 1200.0) == FLOOR_RUNG_MULTILINGUAL


class TestBackendThresholds:
    """Each GPU backend keys to its own measured thresholds (ticket #19)."""

    def test_cuda_keeps_default_only_above_its_own_q8_threshold(self) -> None:
        assert select_ladder_rung(True, 1690.0, "cuda") == DEFAULT_RUNG
        assert select_ladder_rung(True, 1689.0, "cuda") == LOW_MEMORY_RUNG
        # 1644 MiB is enough for Vulkan's q8_0 but not for CUDA's.
        assert select_ladder_rung(True, 1644.0, "cuda") == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, 1644.0, "vulkan") == DEFAULT_RUNG

    def test_cuda_q5_threshold_is_its_own(self) -> None:
        assert select_ladder_rung(True, 1404.0, "cuda") == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, 1403.0, "cuda") == FLOOR_RUNG_ENGLISH
        assert select_ladder_rung(True, 1359.0, "vulkan") == LOW_MEMORY_RUNG
        assert select_ladder_rung(True, 1359.0, "cuda") == FLOOR_RUNG_ENGLISH

    def test_cpu_and_unknown_backends_use_the_vulkan_thresholds(self) -> None:
        for backend in ("cpu", "unknown"):
            assert select_ladder_rung(True, 1644.0, backend) == DEFAULT_RUNG
            assert select_ladder_rung(True, 1359.0, backend) == LOW_MEMORY_RUNG
            assert select_ladder_rung(True, 1358.0, backend) == FLOOR_RUNG_ENGLISH

    def test_default_backend_is_vulkan(self) -> None:
        assert select_ladder_rung(True, 1690.0) == DEFAULT_RUNG
        assert select_ladder_rung(True, 1644.0) == DEFAULT_RUNG

    def test_resolve_passes_the_backend_through(self) -> None:
        assert resolve_startup_rung("", True, 1644.0, "cuda") == LOW_MEMORY_RUNG
        assert resolve_startup_rung("", True, 1644.0, "vulkan") == DEFAULT_RUNG
        assert resolve_startup_rung("tiny", True, 0.0, "cuda") == "tiny"


class TestResolveStartupRung:
    """A manual pin wins verbatim; only unpinned loads use the ladder."""

    def test_pinned_variant_never_overridden(self) -> None:
        assert resolve_startup_rung("tiny", True, 99999.0) == "tiny"
        assert resolve_startup_rung(LOW_MEMORY_RUNG, False, 0.0) == LOW_MEMORY_RUNG
        assert resolve_startup_rung(FLOOR_RUNG_ENGLISH, True, 99999.0) == FLOOR_RUNG_ENGLISH

    def test_unpinned_uses_ladder(self) -> None:
        assert resolve_startup_rung("", True, 99999.0) == DEFAULT_RUNG
        assert resolve_startup_rung("", False, 0.0) == FLOOR_RUNG_MULTILINGUAL


class TestNextRungDown:
    """One step down per failure; the floor fails loudly."""

    def test_default_steps_to_low_memory(self) -> None:
        assert next_rung_down(DEFAULT_RUNG, True) == LOW_MEMORY_RUNG
        assert next_rung_down(DEFAULT_RUNG, False) == LOW_MEMORY_RUNG

    def test_low_memory_steps_to_floor_by_language(self) -> None:
        assert next_rung_down(LOW_MEMORY_RUNG, True) == FLOOR_RUNG_ENGLISH
        assert next_rung_down(LOW_MEMORY_RUNG, False) == FLOOR_RUNG_MULTILINGUAL

    def test_floor_has_no_step_down(self) -> None:
        assert next_rung_down(FLOOR_RUNG_ENGLISH, True) is None
        assert next_rung_down(FLOOR_RUNG_MULTILINGUAL, False) is None

    def test_non_ladder_rungs_have_no_step_down(self) -> None:
        assert next_rung_down("tiny", True) is None
        assert next_rung_down("medium", False) is None


class TestCudaToolkitReadiness:
    """Driver alone is not a toolkit; arch must match too."""

    def test_complete_matching_toolkit_is_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(True, True, (8, 6)) is True
        assert cuda_toolkit_ready_from_probe(True, True, (7, 0)) is True

    def test_driver_without_toolkit_is_not_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(True, False, (8, 6)) is False

    def test_no_gpu_is_not_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(False, True, (8, 6)) is False
        assert cuda_toolkit_ready_from_probe(False, False, None) is False

    def test_old_architecture_is_not_ready(self) -> None:
        assert cuda_toolkit_ready_from_probe(True, True, (6, 1)) is False
        assert cuda_toolkit_ready_from_probe(True, True, None) is False


NVCC_VERSION_OUTPUT = "Cuda compilation tools, release 13.2, V13.2.51\n"
NVIDIA_SMI_OUTPUT = "3711, 8.6\n"


def _write_toolkit(
    root: Path,
    *,
    nvcc: bool = True,
    header: bool = True,
    libcudart: bool = True,
    target: str = "x86_64-linux",
) -> Path:
    """Lay out a CUDA toolkit tree shaped like the NVIDIA runfile install."""
    if nvcc:
        (root / "bin").mkdir(parents=True, exist_ok=True)
        binary = root / "bin" / "nvcc"
        binary.write_text("#!/bin/sh\n")
        binary.chmod(0o755)
    if header:
        (root / "include").mkdir(parents=True, exist_ok=True)
        (root / "include" / "cuda_runtime.h").write_text("/* header */\n")
    if libcudart:
        lib = root / "targets" / target / "lib"
        lib.mkdir(parents=True, exist_ok=True)
        (lib / "libcudart.so.13").write_text("")
        if target == "x86_64-linux":
            (root / "lib64").symlink_to(lib, target_is_directory=True)
    return root


def _probe_stub(nvcc_output: str | None):
    """Stand in for ``_run_probe``: nvidia-smi sees an RTX-class GPU."""

    def probe(argv, timeout=10):
        if argv[0] == "nvidia-smi":
            return NVIDIA_SMI_OUTPUT
        if argv[0] == "nvcc":
            return nvcc_output
        return None

    return probe


@pytest.fixture
def no_env_toolkit(tmp_path: Path):
    """Isolate detection from the host: no CUDA env vars, empty /usr/local."""
    usr_local = tmp_path / "usr-local"
    usr_local.mkdir()
    with (
        patch.dict(os.environ, {}, clear=True),
        patch.object(model_ladder, "_CUDA_USR_LOCAL_ROOT", usr_local),
        patch.object(model_ladder.platform, "machine", return_value="x86_64"),
    ):
        yield usr_local


class TestCudaToolkitDetection:
    """A toolkit counts from PATH, CUDA_HOME, CUDA_PATH or /usr/local/cuda*/."""

    def test_toolkit_found_via_path(self, no_env_toolkit: Path) -> None:
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(NVCC_VERSION_OUTPUT)):
            assert model_ladder.cuda_toolkit_ready() is True

    @pytest.mark.parametrize("variable", ["CUDA_HOME", "CUDA_PATH"])
    def test_toolkit_found_via_cuda_env_root(
        self, tmp_path: Path, no_env_toolkit: Path, variable: str
    ) -> None:
        root = _write_toolkit(tmp_path / "opt" / "cuda")
        with (
            patch.dict(os.environ, {variable: str(root)}, clear=True),
            patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)),
        ):
            assert model_ladder.cuda_toolkit_ready() is True

    def test_toolkit_found_only_under_usr_local_cuda(self, no_env_toolkit: Path) -> None:
        """The real target: /usr/local/cuda -> cuda-13.2 with nvcc off PATH."""
        real = _write_toolkit(no_env_toolkit / "cuda-13.2")
        (no_env_toolkit / "cuda").symlink_to(real, target_is_directory=True)
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is True

    def test_no_toolkit_anywhere_is_not_ready(self, no_env_toolkit: Path) -> None:
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is False

    @pytest.mark.parametrize(
        "layout",
        [
            {"nvcc": False},
            {"header": False},
            {"libcudart": False},
            {"target": "sbsa-linux"},
        ],
        ids=["no-nvcc", "no-cuda-runtime-header", "no-libcudart", "other-architecture"],
    )
    def test_incomplete_toolkit_is_rejected(
        self, tmp_path: Path, no_env_toolkit: Path, layout: dict
    ) -> None:
        root = _write_toolkit(tmp_path / "partial", **layout)
        with (
            patch.dict(os.environ, {"CUDA_HOME": str(root)}, clear=True),
            patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)),
        ):
            assert model_ladder.cuda_toolkit_ready() is False

    def test_driver_without_toolkit_stays_not_ready_with_partial_tree(
        self, no_env_toolkit: Path
    ) -> None:
        """A /usr/local/cuda with headers only is not a toolkit."""
        _write_toolkit(no_env_toolkit / "cuda-13.2", nvcc=False, libcudart=False)
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is False


class TestBuildAwareBackendOrder:
    """The order names only backends the installed build compiled in."""

    def test_vulkan_only_build_never_picks_cuda_even_with_toolkit(
        self, no_env_toolkit: Path
    ) -> None:
        _write_toolkit(no_env_toolkit / "cuda-13.2")
        (no_env_toolkit / "cuda").symlink_to(no_env_toolkit / "cuda-13.2")
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert model_ladder.cuda_toolkit_ready() is True
            assert preferred_backend_order("vulkan") == ("vulkan", "cpu")

    def test_cuda_build_picks_cuda_with_nvcc_off_path(self, no_env_toolkit: Path) -> None:
        _write_toolkit(no_env_toolkit / "cuda-13.2")
        (no_env_toolkit / "cuda").symlink_to(no_env_toolkit / "cuda-13.2")
        with patch.object(model_ladder, "_run_probe", side_effect=_probe_stub(None)):
            assert preferred_backend_order("cuda") == ("cuda", "cpu")

    def test_cpu_build_is_cpu_only(self) -> None:
        assert preferred_backend_order("cpu") == ("cpu",)


class TestPreferredBackendOrder:
    """CUDA ranks first only with a complete toolkit."""

    def test_cuda_first_with_toolkit(self) -> None:
        with patch.object(model_ladder, "cuda_toolkit_ready", return_value=True):
            assert preferred_backend_order() == ("cuda", "vulkan", "cpu")

    def test_vulkan_then_cpu_without_toolkit(self) -> None:
        with patch.object(model_ladder, "cuda_toolkit_ready", return_value=False):
            assert preferred_backend_order() == ("vulkan", "cpu")


def _make_manager(**kwargs):
    """Create a whisper.cpp manager without running any initializer."""
    with (
        patch.object(SpeechRecognitionManager, "_init_vosk"),
        patch.object(SpeechRecognitionManager, "_init_whisper"),
        patch.object(SpeechRecognitionManager, "_init_whispercpp"),
    ):
        return SpeechRecognitionManager(
            engine="whisper_cpp",
            model_size="small",
            language="en-us",
            defer_download=True,
            **kwargs,
        )


class TestStepDownRetry:
    """The loader steps one rung down and retries once, then surfaces."""

    def test_first_failure_steps_down_and_retries_once(self) -> None:
        manager = _make_manager()
        manager.model_size = DEFAULT_RUNG
        calls = []

        def fail_once() -> None:
            calls.append(manager.model_size)
            if len(calls) == 1:
                raise RuntimeError("CUDA out of memory")

        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", side_effect=fail_once):
            manager._init_whispercpp()

        assert calls == [DEFAULT_RUNG, LOW_MEMORY_RUNG]
        assert manager.model_size == LOW_MEMORY_RUNG

    def test_second_failure_surfaces(self) -> None:
        manager = _make_manager()
        manager.model_size = DEFAULT_RUNG
        with patch.object(
            SpeechRecognitionManager, "_init_whispercpp_once", side_effect=RuntimeError("boom")
        ):
            try:
                manager._init_whispercpp()
            except RuntimeError:
                pass
            else:
                raise AssertionError("second failure must surface")
        assert manager.model_size == LOW_MEMORY_RUNG

    def test_manual_pin_never_steps_down(self) -> None:
        manager = _make_manager(whispercpp_variant_pinned=True)
        manager.model_size = DEFAULT_RUNG
        loader = MagicMock(side_effect=RuntimeError("CUDA out of memory"))
        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", loader):
            try:
                manager._init_whispercpp()
            except RuntimeError:
                pass
            else:
                raise AssertionError("manual pick must fail loudly")
        assert loader.call_count == 1
        assert manager.model_size == DEFAULT_RUNG

    def test_floor_rung_fails_loudly(self) -> None:
        manager = _make_manager()
        manager.model_size = FLOOR_RUNG_ENGLISH
        loader = MagicMock(side_effect=OSError("no space"))
        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", loader):
            try:
                manager._init_whispercpp()
            except OSError:
                pass
            else:
                raise AssertionError("floor failure must surface")
        assert loader.call_count == 1

    def test_non_ladder_model_does_not_step_down(self) -> None:
        manager = _make_manager()
        manager.model_size = "medium"
        loader = MagicMock(side_effect=RuntimeError("boom"))
        with patch.object(SpeechRecognitionManager, "_init_whispercpp_once", loader):
            try:
                manager._init_whispercpp()
            except RuntimeError:
                pass
            else:
                raise AssertionError("non-ladder failure must surface")
        assert loader.call_count == 1


def _probe_memory(free_mib: float, backend: str = "cpu"):
    """Stub the startup probes: free memory of the installed build, and its backend."""
    stack = ExitStack()
    stack.enter_context(
        patch(
            "vocalinux.speech_recognition.model_ladder.free_memory_for_backend",
            return_value=(free_mib, "RAM"),
        )
    )
    stack.enter_context(
        patch(
            "vocalinux.speech_recognition.recognition_manager.detect_pywhispercpp_gpu_backend",
            return_value=backend,
        )
    )
    return stack


class TestStartupRungSelection:
    """Unpinned defaults resolve through the ladder; explicit picks do not."""

    def _manager_with(self, saved_size, pinned_variant, language="en-us"):
        from vocalinux.ui.config_manager import ConfigManager

        manager = ConfigManager.__new__(ConfigManager)
        manager.config = {
            "speech_recognition": {
                "whisper_cpp_model_size": saved_size,
                "whisper_cpp_model_variant": pinned_variant,
                "language": language,
            }
        }
        return manager

    def test_unpinned_default_selects_rung_by_memory(self) -> None:
        manager = self._manager_with("tiny", "", "en-us")
        with _probe_memory(99999.0):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == DEFAULT_RUNG

    def test_unpinned_default_low_memory_selects_q5(self) -> None:
        manager = self._manager_with("tiny", "", "en-us")
        with _probe_memory(PREFER_Q5_0_ABOVE_MIB):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == LOW_MEMORY_RUNG

    def test_unpinned_default_tight_memory_selects_floor(self) -> None:
        manager = self._manager_with("tiny", "", "pl")
        with _probe_memory(100.0):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == FLOOR_RUNG_MULTILINGUAL

    def test_pinned_variant_never_uses_ladder(self) -> None:
        manager = self._manager_with("tiny", LOW_MEMORY_RUNG, "en-us")
        with _probe_memory(99999.0):
            assert manager.get_model_size_for_engine("whisper_cpp") == LOW_MEMORY_RUNG

    def test_explicit_size_keeps_legacy_derivation(self) -> None:
        manager = self._manager_with("medium", "", "en-us")
        with _probe_memory(99999.0):
            with patch(
                "vocalinux.utils.whispercpp_model_info.is_model_downloaded", return_value=False
            ):
                assert manager.get_model_size_for_engine("whisper_cpp") == "medium.en"


class TestFreeMemoryForBackend:
    """The memory reading follows the installed build (spec #16, ticket #20)."""

    def test_vulkan_build_reads_discrete_vram_when_present(self) -> None:
        with (
            patch.object(model_ladder, "free_vram_mib", return_value=2600.0),
            patch.object(model_ladder, "free_system_ram_mib", return_value=9000.0),
        ):
            assert model_ladder.free_memory_for_backend("vulkan") == (2600.0, "VRAM")

    def test_cuda_build_reads_discrete_vram_when_present(self) -> None:
        with patch.object(model_ladder, "free_vram_mib", return_value=1500.0):
            assert model_ladder.free_memory_for_backend("cuda") == (1500.0, "VRAM")

    def test_gpu_build_without_an_nvidia_reading_reads_ram(self) -> None:
        """Integrated graphics: no discrete VRAM reading, so system RAM decides."""
        with (
            patch.object(model_ladder, "free_vram_mib", return_value=None),
            patch.object(model_ladder, "free_system_ram_mib", return_value=6000.0),
        ):
            assert model_ladder.free_memory_for_backend("vulkan") == (6000.0, "RAM")

    def test_cpu_build_never_reads_vram(self) -> None:
        vram = MagicMock(return_value=2600.0)
        with (
            patch.object(model_ladder, "free_vram_mib", vram),
            patch.object(model_ladder, "free_system_ram_mib", return_value=1200.0),
        ):
            assert model_ladder.free_memory_for_backend("cpu") == (1200.0, "RAM")
        vram.assert_not_called()