# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Tests for is_vision_model() caching behaviour.

``_vision_detection_cache`` mirrors the ``_audio_detection_cache``
pattern used by ``detect_audio_type()``. These tests verify:

* Repeated calls for the same model hit the cache.
* Different models each trigger their own detection.
* Both True and False results are cached.
* The subprocess path (transformers 5.x models) is cached.
* Exceptions that fall back to False are cached.
"""

import sys
import types as _types
from pathlib import Path
from unittest.mock import patch, MagicMock

import pytest

# sys.path + logger stub — same pattern as the rest of the test suite
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
    sys.path.insert(0, _BACKEND_DIR)

_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)

from utils.models.model_config import (
    is_vision_model,
    _is_vision_model_uncached,
    _vision_detection_cache,
)


# Helpers


@pytest.fixture(autouse = True)
def _clear_vision_cache():
    """Ensure every test starts with a fresh cache."""
    _vision_detection_cache.clear()
    yield
    _vision_detection_cache.clear()


# Cache hit / miss tests


class TestVisionCacheHitMiss:
    """Verify the cache prevents redundant detection calls."""

    @patch("utils.models.model_config._is_vision_model_uncached", return_value = True)
    def test_second_call_uses_cache(self, mock_uncached):
        """Two calls for the same model invoke the uncached fn once."""
        assert is_vision_model("org/my-vlm") is True
        assert is_vision_model("org/my-vlm") is True
        mock_uncached.assert_called_once_with("org/my-vlm", None)

    @patch("utils.models.model_config._is_vision_model_uncached", return_value = False)
    def test_different_models_each_detected(self, mock_uncached):
        """Different model names should each trigger detection."""
        is_vision_model("model-a")
        is_vision_model("model-b")
        assert mock_uncached.call_count == 2

    @patch("utils.models.model_config._is_vision_model_uncached", return_value = True)
    def test_cache_returns_correct_value(self, mock_uncached):
        """The cached value must match what _is_vision_model_uncached returned."""
        first = is_vision_model("org/vlm")
        second = is_vision_model("org/vlm")
        assert first is True
        assert second is True


class TestVisionCacheStoresFalse:
    """Non-VLM results (False) must also be cached to avoid re-detection."""

    @patch("utils.models.model_config._is_vision_model_uncached", return_value = False)
    def test_false_result_cached(self, mock_uncached):
        assert is_vision_model("org/text-only") is False
        assert is_vision_model("org/text-only") is False
        mock_uncached.assert_called_once()
        assert _vision_detection_cache[("org/text-only", None)] is False


# Subprocess path (transformers 5.x) caching


class TestVisionCacheSubprocessPath:
    """transformers 5.x models go through _is_vision_model_subprocess.
    The cache should spawn the subprocess at most once per model per
    process."""

    @patch("utils.models.model_config._is_vision_model_subprocess", return_value = True)
    @patch("utils.transformers_version.needs_transformers_5", return_value = True)
    def test_subprocess_called_once_with_cache(self, mock_needs_t5, mock_subprocess):
        """Subprocess fires only on the first call; second is cached."""
        # First call: uncached → subprocess
        assert is_vision_model("unsloth/Qwen3.5-2B") is True
        # Second call: cache hit, no subprocess
        assert is_vision_model("unsloth/Qwen3.5-2B") is True

        mock_subprocess.assert_called_once()
        assert _vision_detection_cache[("unsloth/Qwen3.5-2B", None)] is True

    @patch("utils.models.model_config._raw_config_has_vision_config", return_value = True)
    @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None)
    @patch("utils.transformers_version.needs_transformers_5", return_value = True)
    def test_subprocess_none_falls_back_to_raw_vision_config(
        self, mock_needs_t5, mock_subprocess, mock_raw_config
    ):
        assert is_vision_model("unsloth/gemma-4-E4B-it") is True
        assert is_vision_model("unsloth/gemma-4-E4B-it") is True

        mock_subprocess.assert_called_once()
        mock_raw_config.assert_called_once_with("unsloth/gemma-4-E4B-it", hf_token = None)


# Exception handling — cache the False fallback


class TestVisionCacheOnException:
    """On exception, _is_vision_model_uncached distinguishes permanent
    failures (cached as False) from transient ones (returned as None,
    not cached, so the next call retries). Verify both contracts."""

    @patch(
        "utils.models.model_config.load_model_config",
        side_effect = ValueError("bad config"),
    )
    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    def test_permanent_exception_result_cached(self, mock_needs_t5, mock_load_config):
        """A permanent failure (ValueError / RepositoryNotFoundError /
        GatedRepoError / JSONDecodeError) is caught, returns False, and
        that False is cached so subsequent calls don't retry. ValueError
        stands in as the simplest cacheable exception type."""
        # First call raises -> False; second is a cache hit.
        assert is_vision_model("broken/model") is False
        assert is_vision_model("broken/model") is False
        mock_load_config.assert_called_once()

    @patch(
        "utils.models.model_config.load_model_config",
        side_effect = OSError("network down"),
    )
    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    def test_transient_exception_not_cached(self, mock_needs_t5, mock_load_config):
        """A transient failure (OSError, timeouts) returns None from
        _is_vision_model_uncached, surfaces as False, and is NOT cached
        so the next call retries."""
        # First call: OSError -> False, not cached; second call retries.
        assert is_vision_model("broken/model") is False
        assert is_vision_model("broken/model") is False
        assert mock_load_config.call_count == 2


# Direct detection path (non-transformers-5 models) caching


class TestVisionCacheDirectPath:
    """Models that do NOT need transformers 5.x detect via
    load_model_config directly. The cache must work the same way."""

    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    @patch("utils.models.model_config.load_model_config")
    def test_direct_vlm_detection_cached(self, mock_load_config, mock_needs_t5):
        """A standard VLM detected via architecture suffix should be cached."""
        cfg = MagicMock(spec = [])  # strict: only explicitly set attrs exist
        cfg.model_type = "gemma3"
        cfg.architectures = ["Gemma3ForConditionalGeneration"]
        mock_load_config.return_value = cfg

        assert is_vision_model("google/gemma-3-4b-it") is True
        assert is_vision_model("google/gemma-3-4b-it") is True
        # load_model_config should only be called once
        mock_load_config.assert_called_once()

    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    @patch("utils.models.model_config.load_model_config")
    def test_direct_non_vlm_detection_cached(self, mock_load_config, mock_needs_t5):
        """A standard text model (no VLM indicators) should cache False."""
        cfg = MagicMock(spec = [])  # spec=[] means no attributes at all
        cfg.model_type = "llama"
        cfg.architectures = ["LlamaForCausalLM"]
        mock_load_config.return_value = cfg

        # No VLM suffix, no vision_config, etc.
        assert is_vision_model("meta-llama/Llama-3-8B") is False
        assert is_vision_model("meta-llama/Llama-3-8B") is False
        mock_load_config.assert_called_once()

    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    @patch("utils.models.model_config.load_model_config")
    def test_vision_config_attr_detected_and_cached(self, mock_load_config, mock_needs_t5):
        """Models with vision_config (LLaVA, Qwen2-VL, etc.) should be cached as True."""
        cfg = MagicMock(spec = [])  # strict: only explicitly set attrs exist
        cfg.model_type = "qwen2_vl"
        cfg.architectures = ["Qwen2VLForCausalLM"]  # Doesn't match VLM suffixes
        cfg.vision_config = {"hidden_size": 1024}
        mock_load_config.return_value = cfg

        assert is_vision_model("Qwen/Qwen2-VL-7B") is True
        assert is_vision_model("Qwen/Qwen2-VL-7B") is True
        mock_load_config.assert_called_once()

    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    @patch("utils.models.model_config.load_model_config")
    def test_gemma4_model_type_detected_and_cached(self, mock_load_config, mock_needs_t5):
        cfg = MagicMock(spec = [])
        cfg.model_type = "gemma4"
        cfg.architectures = ["Gemma4ForConditionalGeneration"]
        mock_load_config.return_value = cfg

        assert is_vision_model("google/gemma-4-E4B-it") is True
        assert is_vision_model("google/gemma-4-E4B-it") is True
        mock_load_config.assert_called_once()

    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    @patch("utils.models.model_config.load_model_config")
    def test_gemma4_audio_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5):
        cfg = MagicMock(spec = [])
        cfg.model_type = "gemma4_audio"
        cfg.architectures = ["Gemma4AudioModel"]
        mock_load_config.return_value = cfg

        assert is_vision_model("local/gemma4-audio-encoder") is False
        assert is_vision_model("local/gemma4-audio-encoder") is False
        mock_load_config.assert_called_once()

    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    @patch("utils.models.model_config.load_model_config")
    def test_gemma4_text_subconfig_not_detected_as_vision(self, mock_load_config, mock_needs_t5):
        cfg = MagicMock(spec = [])
        cfg.model_type = "gemma4_text"
        cfg.architectures = ["Gemma4ForCausalLM"]
        mock_load_config.return_value = cfg

        assert is_vision_model("local/gemma-4-text") is False
        assert is_vision_model("local/gemma-4-text") is False
        mock_load_config.assert_called_once()

    @patch("utils.transformers_version.needs_transformers_5", return_value = False)
    @patch("utils.models.model_config.load_model_config")
    def test_audio_model_excluded_and_cached(self, mock_load_config, mock_needs_t5):
        """Audio-only models (csm, whisper) with ForConditionalGeneration
        should be excluded from VLM detection and cached as False."""
        cfg = MagicMock(spec = [])  # strict: only explicitly set attrs exist
        cfg.model_type = "whisper"
        cfg.architectures = ["WhisperForConditionalGeneration"]
        mock_load_config.return_value = cfg

        assert is_vision_model("openai/whisper-large-v3") is False
        assert is_vision_model("openai/whisper-large-v3") is False
        mock_load_config.assert_called_once()


# hf_token handling


class TestVisionCacheTokenHandling:
    """The cache is keyed on (model_name, hf_token). Different tokens
    for the same model trigger separate detections for gated models."""

    @patch("utils.models.model_config._is_vision_model_uncached", return_value = True)
    def test_different_tokens_trigger_new_detection(self, mock_uncached):
        """Different tokens trigger separate detections for gated models
        (e.g. unauthenticated probe → False, then authenticated
        re-check)."""
        assert is_vision_model("gated/model", hf_token = "token-a") is True
        assert is_vision_model("gated/model", hf_token = "token-b") is True
        assert mock_uncached.call_count == 2

    @patch("utils.models.model_config._is_vision_model_uncached", return_value = True)
    def test_same_token_uses_cache(self, mock_uncached):
        """Repeated calls with identical model + token should hit cache."""
        assert is_vision_model("gated/model", hf_token = "token-a") is True
        assert is_vision_model("gated/model", hf_token = "token-a") is True
        mock_uncached.assert_called_once()


# ---------------------------------------------------------------------------
# Direct unit tests for _raw_config_has_vision_config
# ---------------------------------------------------------------------------


import json as _json

from utils.models.model_config import (
    _AUDIO_ONLY_MODEL_TYPES,
    _VISION_CHECK_INLINE_HELPERS,
    _VISION_CHECK_SCRIPT,
    _is_vlm,
    _raw_config_has_vision_config,
)


def _write_config(tmp_path, config):
    (tmp_path / "config.json").write_text(_json.dumps(config))
    return tmp_path


class TestRawConfigVlmDetection:
    """Direct coverage of _raw_config_has_vision_config across the same
    indicator set used by _is_vlm. The cache integration tests above mock
    this function; these exercise its real implementation."""

    def test_truthy_vision_config(self, tmp_path):
        p = _write_config(tmp_path, {"vision_config": {"hidden_size": 1024}})
        assert _raw_config_has_vision_config(str(p)) is True

    def test_empty_vision_config_key(self, tmp_path):
        p = _write_config(tmp_path, {"vision_config": {}})
        assert _raw_config_has_vision_config(str(p)) is True

    def test_arch_suffix_detection(self, tmp_path):
        p = _write_config(
            tmp_path,
            {
                "architectures": ["Gemma4ForConditionalGeneration"],
                "model_type": "gemma4",
            },
        )
        assert _raw_config_has_vision_config(str(p)) is True

    def test_img_processor_key(self, tmp_path):
        p = _write_config(tmp_path, {"img_processor": {"image_size": 336}})
        assert _raw_config_has_vision_config(str(p)) is True

    def test_image_token_index_key(self, tmp_path):
        p = _write_config(tmp_path, {"image_token_index": 32000})
        assert _raw_config_has_vision_config(str(p)) is True

    def test_known_vlm_model_type(self, tmp_path):
        p = _write_config(tmp_path, {"model_type": "gemma4"})
        assert _raw_config_has_vision_config(str(p)) is True

    def test_plain_text_model_returns_false(self, tmp_path):
        p = _write_config(
            tmp_path,
            {"model_type": "llama", "architectures": ["LlamaForCausalLM"]},
        )
        assert _raw_config_has_vision_config(str(p)) is False

    def test_missing_config_returns_none(self, tmp_path):
        assert _raw_config_has_vision_config(str(tmp_path)) is None


# ---------------------------------------------------------------------------
# Self-contained subprocess script (no parent backend imports)
# ---------------------------------------------------------------------------


class TestSubprocessScript:
    def test_does_not_import_parent_module(self):
        assert "from utils.models.model_config" not in _VISION_CHECK_SCRIPT

    def test_inline_is_vlm_executes_correctly(self):
        ns: dict = {}
        exec(_VISION_CHECK_INLINE_HELPERS, ns)
        inline_is_vlm = ns["_is_vlm"]

        class _C:
            def __init__(self, **kw):
                for k, v in kw.items():
                    setattr(self, k, v)

        assert (
            inline_is_vlm(
                _C(
                    model_type = "gemma4",
                    architectures = ["Gemma4ForConditionalGeneration"],
                )
            )
            is True
        )
        assert (
            inline_is_vlm(_C(model_type = "gemma4_text", architectures = ["Gemma4ForCausalLM"]))
            is False
        )
        assert inline_is_vlm(_C(model_type = "llama", architectures = ["LlamaForCausalLM"])) is False


# ---------------------------------------------------------------------------
# Audio-only model exclusion must apply across every detection path
# ---------------------------------------------------------------------------


class TestVlmAudioExclusion:
    """The {csm, whisper} guard previously lived only in the direct caller
    branch. These tests assert it now applies inside _is_vlm, the raw
    fallback, and the inlined subprocess helper too."""

    def test_audio_only_set_canonical(self):
        assert _AUDIO_ONLY_MODEL_TYPES == {"csm", "whisper"}

    def test_is_vlm_excludes_whisper(self):
        cfg = MagicMock(spec = [])
        cfg.model_type = "whisper"
        cfg.architectures = ["WhisperForConditionalGeneration"]
        assert _is_vlm(cfg) is False

    def test_raw_fallback_excludes_whisper(self, tmp_path):
        p = _write_config(
            tmp_path,
            {
                "architectures": ["WhisperForConditionalGeneration"],
                "model_type": "whisper",
            },
        )
        assert _raw_config_has_vision_config(str(p)) is False

    def test_inline_subprocess_helper_excludes_whisper(self):
        ns: dict = {}
        exec(_VISION_CHECK_INLINE_HELPERS, ns)
        cfg = MagicMock(spec = [])
        cfg.model_type = "whisper"
        cfg.architectures = ["WhisperForConditionalGeneration"]
        assert ns["_is_vlm"](cfg) is False

    @patch("utils.models.model_config._is_vision_model_subprocess", return_value = None)
    @patch("utils.transformers_version.needs_transformers_5", return_value = True)
    def test_t5_subprocess_none_falls_back_through_raw_for_whisper(
        self, mock_needs_t5, mock_subprocess, tmp_path
    ):
        _write_config(
            tmp_path,
            {
                "architectures": ["WhisperForConditionalGeneration"],
                "model_type": "whisper",
            },
        )
        assert is_vision_model(str(tmp_path)) is False