# SPDX-License-Identifier: Apache-2.0
"""Tests for the DeepSeek V4 monkey-patch (PR 1192 port)."""

import importlib
import inspect
import sys
from types import SimpleNamespace

import pytest


@pytest.fixture(scope="module")
def applied_patch():
    """Apply the patch once for the whole module. The patch itself is
    idempotent so repeated calls are safe."""
    from omlx.patches.deepseek_v4 import apply_deepseek_v4_patch

    apply_deepseek_v4_patch()
    return True


class TestPatchOrchestration:
    """Top-level apply / idempotency / module registration checks."""

    def test_apply_returns_true_first_time(self):
        from omlx.patches.deepseek_v4 import apply_deepseek_v4_patch, is_applied

        # The patch may have been applied by a previous test run in the
        # same process; force-reset is_applied to validate the flow.
        # The module-level _APPLIED guard means we cannot un-apply, so
        # this test is informational about the *current* state.
        if is_applied():
            assert apply_deepseek_v4_patch() is False
        else:
            assert apply_deepseek_v4_patch() is True
            assert is_applied() is True

    def test_apply_is_idempotent(self, applied_patch):
        from omlx.patches.deepseek_v4 import apply_deepseek_v4_patch

        # After fixture has applied the patch, a second call must return False.
        assert apply_deepseek_v4_patch() is False

    def test_hyper_connection_registered(self, applied_patch):
        assert "mlx_lm.models.hyper_connection" in sys.modules

    def test_deepseek_v4_registered(self, applied_patch):
        assert "mlx_lm.models.deepseek_v4" in sys.modules

    def test_deepseek_v4_module_package(self, applied_patch):
        mod = sys.modules["mlx_lm.models.deepseek_v4"]
        # __package__ must be mlx_lm.models so relative imports inside
        # the loaded file resolve through the real mlx_lm package.
        assert mod.__package__ == "mlx_lm.models"

    def test_deepseek_v4_mtp_alias_registered(self, applied_patch):
        assert (
            sys.modules["mlx_lm.models.deepseek_v4_mtp"]
            is sys.modules["mlx_lm.models.deepseek_v4"]
        )


class TestCacheInjection:
    """PoolingCache / BatchPoolingCache injected into mlx_lm.models.cache."""

    def test_pooling_cache_attribute(self, applied_patch):
        import mlx_lm.models.cache as cache_mod

        assert hasattr(cache_mod, "PoolingCache")
        assert hasattr(cache_mod, "BatchPoolingCache")

    def test_pooling_cache_module_attribute(self, applied_patch):
        from mlx_lm.models.cache import BatchPoolingCache, PoolingCache

        # The injected classes claim to live in mlx_lm.models.cache so
        # any introspection (e.g. type(c).__module__) sees the right name.
        assert PoolingCache.__module__ == "mlx_lm.models.cache"
        assert BatchPoolingCache.__module__ == "mlx_lm.models.cache"

    def test_pooling_cache_instantiation(self, applied_patch):
        from mlx_lm.models.cache import PoolingCache

        cache = PoolingCache(ratio=4)
        assert cache.ratio == 4
        assert cache.empty()
        assert cache.size() == 0
        assert cache.offset == 0


class TestUtilsPatch:
    """mlx_lm.utils.load_model + _load_safetensors + SAFETENSORS_DTYPE_FALLBACKS."""

    def test_load_model_replaced(self, applied_patch):
        import mlx_lm.utils as utils_mod

        # The replaced function carries our docstring marker via its
        # bound name; just check it's not the upstream one by virtue of
        # the new attributes around it.
        assert hasattr(utils_mod, "_load_safetensors")
        assert hasattr(utils_mod, "SAFETENSORS_DTYPE_FALLBACKS")

    def test_dtype_fallback_map(self, applied_patch):
        import mlx_lm.utils as utils_mod

        assert utils_mod.SAFETENSORS_DTYPE_FALLBACKS == {"F8_E8M0": "U8"}

    def test_load_safetensors_passthrough_for_normal_dtype(
        self, applied_patch, tmp_path
    ):
        """A safetensors file with a standard dtype must round-trip
        through _load_safetensors unchanged (no header rewrite)."""
        import mlx.core as mx
        from mlx_lm.utils import _load_safetensors

        path = tmp_path / "model.safetensors"
        data = {"x": mx.zeros((4, 4), dtype=mx.float32)}
        mx.save_safetensors(str(path), data)
        loaded = _load_safetensors(str(path))
        assert "x" in loaded
        assert loaded["x"].shape == (4, 4)


class TestGeneratePatch:
    """mlx_lm.generate._make_cache replaced."""

    def test_make_cache_replaced(self, applied_patch):
        gen_mod = importlib.import_module("mlx_lm.generate")

        assert hasattr(gen_mod, "_make_cache")
        # Source must include PoolingCache → BatchPoolingCache branch.
        # We can't easily compare functions, so just verify the new
        # behavior: passing a model with a PoolingCache in make_cache
        # produces a BatchPoolingCache.
        from mlx_lm.models.cache import BatchPoolingCache, PoolingCache

        class FakeModel:
            def __init__(self):
                self.layers = [None]

            def make_cache(self):
                return [PoolingCache(ratio=4)]

        result = gen_mod._make_cache(FakeModel(), [0], None)
        assert len(result) == 1
        assert isinstance(result[0], BatchPoolingCache)


class TestTokenizerPatch:
    """mlx_lm.tokenizer_utils.AutoTokenizer wrapped with deepseek_v4 fallback."""

    def test_autotokenizer_wrapped(self, applied_patch):
        import mlx_lm.tokenizer_utils as tu

        # Wrapped class still exposes from_pretrained.
        assert hasattr(tu.AutoTokenizer, "from_pretrained")
        # Class name preserved for any introspection.
        assert tu.AutoTokenizer.__name__ == "AutoTokenizer"

    def test_passthrough_on_success(self, applied_patch):
        """When upstream AutoTokenizer.from_pretrained succeeds, the wrapper
        must return its result unmodified — no fallback path taken."""
        from unittest.mock import patch as mock_patch

        from omlx.patches.deepseek_v4 import tokenizer_patch

        sentinel = object()

        class _FakeUpstream:
            calls = []

            @staticmethod
            def from_pretrained(model_path, *args, **kwargs):
                _FakeUpstream.calls.append((model_path, args, kwargs))
                return sentinel

        with mock_patch("transformers.AutoTokenizer", _FakeUpstream):
            wrapper = tokenizer_patch._build_wrapper()
            result = wrapper.from_pretrained("/fake/path", trust_remote_code=True)

        assert result is sentinel
        assert len(_FakeUpstream.calls) == 1
        # Fallback never injected its own config kwarg.
        assert "config" not in _FakeUpstream.calls[0][2]

    def test_fallback_on_max_position_embeddings_error(self, applied_patch):
        """The exact AttributeError that transformers raises when it cannot
        recognize deepseek_v4 must trigger a retry with PreTrainedConfig()."""
        import pytest as _pytest
        from unittest.mock import patch as mock_patch

        from omlx.patches.deepseek_v4 import tokenizer_patch

        class _FakeUpstream:
            calls = []

            @staticmethod
            def from_pretrained(model_path, *args, **kwargs):
                _FakeUpstream.calls.append((model_path, args, kwargs))
                if "config" in kwargs:
                    return "FALLBACK_OK"
                raise AttributeError(
                    "'PreTrainedConfig' object has no attribute "
                    "'max_position_embeddings'"
                )

        with mock_patch("transformers.AutoTokenizer", _FakeUpstream):
            wrapper = tokenizer_patch._build_wrapper()
            with _pytest.warns(
                RuntimeWarning, match="Falling back to generic tokenizer config"
            ):
                result = wrapper.from_pretrained("/fake/path")

        assert result == "FALLBACK_OK"
        assert len(_FakeUpstream.calls) == 2
        # Second call must inject config=PreTrainedConfig().
        assert "config" in _FakeUpstream.calls[1][2]

    def test_fallback_on_deepseek_v4_value_error(self, applied_patch):
        """ValueError mentioning deepseek_v4 also triggers fallback."""
        import pytest as _pytest
        from unittest.mock import patch as mock_patch

        from omlx.patches.deepseek_v4 import tokenizer_patch

        class _FakeUpstream:
            calls = []

            @staticmethod
            def from_pretrained(model_path, *args, **kwargs):
                _FakeUpstream.calls.append((model_path, args, kwargs))
                if "config" in kwargs:
                    return "FALLBACK_OK"
                raise ValueError("Unrecognized configuration class for deepseek_v4")

        with mock_patch("transformers.AutoTokenizer", _FakeUpstream):
            wrapper = tokenizer_patch._build_wrapper()
            with _pytest.warns(
                RuntimeWarning, match="Falling back to generic tokenizer config"
            ):
                result = wrapper.from_pretrained("/fake/path")

        assert result == "FALLBACK_OK"
        assert len(_FakeUpstream.calls) == 2

    def test_unrelated_error_reraises(self, applied_patch):
        """Errors outside the deepseek_v4 / max_position_embeddings signature
        must NOT be swallowed."""
        from unittest.mock import patch as mock_patch

        import pytest as _pytest

        from omlx.patches.deepseek_v4 import tokenizer_patch

        class _FakeUpstream:
            @staticmethod
            def from_pretrained(*args, **kwargs):
                raise ValueError("totally unrelated error")

        with mock_patch("transformers.AutoTokenizer", _FakeUpstream):
            wrapper = tokenizer_patch._build_wrapper()
            with _pytest.raises(ValueError, match="totally unrelated"):
                wrapper.from_pretrained("/fake/path")

    def test_explicit_config_skips_fallback(self, applied_patch):
        """If the caller already passed config=, we must not override it
        even when the inner call raises a matching error."""
        from unittest.mock import patch as mock_patch

        import pytest as _pytest

        from omlx.patches.deepseek_v4 import tokenizer_patch

        class _FakeUpstream:
            @staticmethod
            def from_pretrained(*args, **kwargs):
                # Caller-provided config is in kwargs; we still raise the
                # max_position_embeddings error to verify the wrapper does
                # not silently retry.
                raise AttributeError(
                    "'PreTrainedConfig' object has no attribute "
                    "'max_position_embeddings'"
                )

        with mock_patch("transformers.AutoTokenizer", _FakeUpstream):
            wrapper = tokenizer_patch._build_wrapper()
            with _pytest.raises(AttributeError, match="max_position_embeddings"):
                wrapper.from_pretrained("/fake/path", config="caller_supplied")

    def test_class_attribute_forwarding(self, applied_patch):
        """Class-level attribute access (e.g. AutoTokenizer.register) must
        forward to the upstream class so mlx-lm's NewlineTokenizer
        registration still works."""
        import mlx_lm.tokenizer_utils as tu
        from transformers import AutoTokenizer as upstream_at

        # register is an upstream classmethod — wrapped class must expose it.
        assert tu.AutoTokenizer.register is upstream_at.register


class TestDSMLToolParser:
    """tool_parser_v4 — DSML invoke / parameter grammar parsing."""

    def test_single_invoke_typed_args(self, applied_patch):
        from omlx.patches.deepseek_v4 import tool_parser_v4 as tp

        text = (
            '<｜DSML｜invoke name="get_weather">\n'
            '<｜DSML｜parameter name="city" string="true">Seoul</｜DSML｜parameter>\n'
            '<｜DSML｜parameter name="days" string="false">7</｜DSML｜parameter>\n'
            '<｜DSML｜parameter name="imperial" string="false">false</｜DSML｜parameter>\n'
            "</｜DSML｜invoke>"
        )
        result = tp.parse_tool_call(text)
        assert result["name"] == "get_weather"
        assert result["arguments"] == {
            "city": "Seoul",
            "days": 7,
            "imperial": False,
        }

    def test_multiple_invokes_returns_list(self, applied_patch):
        from omlx.patches.deepseek_v4 import tool_parser_v4 as tp

        text = (
            '<｜DSML｜invoke name="a">'
            '<｜DSML｜parameter name="x" string="false">1</｜DSML｜parameter>'
            "</｜DSML｜invoke>\n"
            '<｜DSML｜invoke name="b">'
            '<｜DSML｜parameter name="y" string="true">hello</｜DSML｜parameter>'
            "</｜DSML｜invoke>"
        )
        result = tp.parse_tool_call(text)
        assert isinstance(result, list)
        assert len(result) == 2
        assert result[0] == {"name": "a", "arguments": {"x": 1}}
        assert result[1] == {"name": "b", "arguments": {"y": "hello"}}

    def test_object_and_array_parameters(self, applied_patch):
        from omlx.patches.deepseek_v4 import tool_parser_v4 as tp

        text = (
            '<｜DSML｜invoke name="search">\n'
            '<｜DSML｜parameter name="filters" string="false">'
            '{"category": "books", "min_price": 10}'
            "</｜DSML｜parameter>\n"
            '<｜DSML｜parameter name="ids" string="false">[1, 2, 3]</｜DSML｜parameter>\n'
            "</｜DSML｜invoke>"
        )
        result = tp.parse_tool_call(text)
        assert result["arguments"]["filters"] == {"category": "books", "min_price": 10}
        assert result["arguments"]["ids"] == [1, 2, 3]

    def test_no_invoke_raises(self, applied_patch):
        import pytest as _pytest

        from omlx.patches.deepseek_v4 import tool_parser_v4 as tp

        with _pytest.raises(ValueError, match="No.*invoke.*block"):
            tp.parse_tool_call("just some plain text without DSML markup")

    def test_outer_markers_exposed(self, applied_patch):
        from omlx.patches.deepseek_v4 import tool_parser_v4 as tp

        # mlx-lm reads these as module attributes for stream detection.
        assert tp.tool_call_start == "<｜DSML｜tool_calls>"
        assert tp.tool_call_end == "</｜DSML｜tool_calls>"


class TestChatTemplateV4:
    """chat_template_v4 — DSML system prompt + tool_calls render."""

    def test_outer_marker_uses_tool_calls_not_function_calls(self, applied_patch):
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct

        # vllm's DeepSeekV4ToolParser overrides only the outer marker
        # name (tool_calls vs V3.2's function_calls). Verify our copy
        # made that one edit.
        assert "function_calls" not in ct.tool_calls_template
        assert "tool_calls" in ct.tool_calls_template
        assert "function_calls" not in ct.TOOLS_SYSTEM_TEMPLATE

    def test_inner_grammar_unchanged_from_v32(self, applied_patch):
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct

        # Inner markers must still be invoke / parameter — V4 reuses V3.2's
        # invoke/parameter grammar.
        assert "invoke" in ct.tool_call_template
        assert "parameter" in ct.encode_arguments_to_dsml(
            {"name": "x", "arguments": '{"k": "v"}'}
        )

    def test_round_trip_encode_then_parse(self, applied_patch):
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct
        from omlx.patches.deepseek_v4 import tool_parser_v4 as tp

        encoded_args = ct.encode_arguments_to_dsml(
            {"name": "f", "arguments": '{"a": 1, "b": "hi", "c": [1, 2]}'}
        )
        invoke = ct.tool_call_template.format(
            dsml_token=ct.dsml_token, name="f", arguments=encoded_args
        )
        block = ct.tool_calls_template.format(
            dsml_token=ct.dsml_token, tool_calls=invoke
        )
        # Strip the outer markers as TokenizerWrapper would.
        inner = (
            block.replace(tp.tool_call_start, "").replace(tp.tool_call_end, "").strip()
        )
        parsed = tp.parse_tool_call(inner)
        assert parsed == {"name": "f", "arguments": {"a": 1, "b": "hi", "c": [1, 2]}}

    def test_user_only_request_with_tools_injects_dsml(self, applied_patch):
        """User-only message + tools must still emit the DSML tools block.

        Regression guard for the case where a Claude Code or OpenAI client
        passes ``tools`` without a system message. ``render_message`` only
        injects tools on system / developer roles, so ``encode_messages``
        synthesises an empty system message up front when the first
        message is a plain user. Without this fix the rendered prompt
        omits the ``<functions>`` schema entirely and the model never
        emits a tool_calls block.
        """
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct

        tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "Get the weather",
                    "parameters": {
                        "type": "object",
                        "properties": {"location": {"type": "string"}},
                        "required": ["location"],
                    },
                },
            }
        ]
        prompt = ct.apply_chat_template(
            [{"role": "user", "content": "Weather in Seoul?"}],
            tools=tools,
            add_generation_prompt=True,
        )
        assert "<functions>" in prompt
        assert "get_weather" in prompt
        assert ct.dsml_token in prompt

    def test_system_user_request_with_tools_unchanged(self, applied_patch):
        """When a system message is already present, the synthetic prepend
        path must not fire — the rendered prompt keeps the original system
        content verbatim and only injects the tools schema once.
        """
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct

        tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "Get the weather",
                    "parameters": {
                        "type": "object",
                        "properties": {"location": {"type": "string"}},
                        "required": ["location"],
                    },
                },
            }
        ]
        prompt = ct.apply_chat_template(
            [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": "Weather in Seoul?"},
            ],
            tools=tools,
            add_generation_prompt=True,
        )
        assert "You are a helpful assistant." in prompt
        # Only one tools block — no double-injection from synthetic prepend.
        assert prompt.count("<functions>") == 1

    def test_user_only_no_tools_no_prepend(self, applied_patch):
        """No tools → no synthetic system. Plain user-only request renders
        with just the BOS + user wrapper, matching V3.2 baseline."""
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct

        prompt = ct.apply_chat_template(
            [{"role": "user", "content": "Hi"}],
            add_generation_prompt=True,
        )
        assert "<functions>" not in prompt
        assert "## Tools" not in prompt

    def test_encode_arguments_accepts_dict(self, applied_patch):
        """Anthropic /v1/messages history stores tool_call arguments as
        a dict (anthropic_utils.py decodes the input before saving).
        encode_arguments_to_dsml must accept that shape — not just the
        OpenAI JSON-string convention — so multi-turn renders don't
        raise TypeError when the assistant history is from Claude Code.
        """
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct

        encoded = ct.encode_arguments_to_dsml(
            {"name": "f", "arguments": {"location": "Seoul", "n": 3}}
        )
        assert 'name="location"' in encoded and "Seoul" in encoded
        assert 'name="n"' in encoded and ">3<" in encoded
        # string="true" for string params, "false" for non-string.
        assert 'string="true"' in encoded
        assert 'string="false"' in encoded

    def test_assistant_tool_call_dict_arguments_round_trip(self, applied_patch):
        """End-to-end multi-turn: assistant message history contains a
        tool_use whose arguments came in as dict (Anthropic shape). The
        rendered prompt must include the assistant's prior tool_call
        block in DSML form so the model can continue the conversation
        coherently.
        """
        from omlx.patches.deepseek_v4 import chat_template_v4 as ct

        messages = [
            {"role": "user", "content": "Weather in Seoul?"},
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {
                            "name": "get_weather",
                            "arguments": {"location": "Seoul"},
                        },
                    }
                ],
            },
            {"role": "tool", "content": "sunny, 22C"},
        ]
        tools = [
            {
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "Get the weather",
                    "parameters": {
                        "type": "object",
                        "properties": {"location": {"type": "string"}},
                        "required": ["location"],
                    },
                },
            }
        ]
        prompt = ct.apply_chat_template(
            messages, tools=tools, add_generation_prompt=True
        )
        assert "<｜DSML｜tool_calls>" in prompt
        assert 'invoke name="get_weather"' in prompt
        assert "Seoul" in prompt
        assert "sunny, 22C" in prompt


class TestChatTemplateModuleRegistration:
    """sys.modules registration so mlx-lm's importlib path picks up our types."""

    def test_chat_template_module_registered(self, applied_patch):
        import sys

        assert "mlx_lm.chat_templates.deepseek_v4" in sys.modules
        mod = sys.modules["mlx_lm.chat_templates.deepseek_v4"]
        assert hasattr(mod, "apply_chat_template")

    def test_tool_parser_module_registered(self, applied_patch):
        import sys

        assert "mlx_lm.tool_parsers.deepseek_v4" in sys.modules
        mod = sys.modules["mlx_lm.tool_parsers.deepseek_v4"]
        assert hasattr(mod, "parse_tool_call")
        assert mod.tool_call_start == "<｜DSML｜tool_calls>"
        assert mod.tool_call_end == "</｜DSML｜tool_calls>"


class TestModelClassResolution:
    """mlx_lm.utils._get_classes resolves deepseek_v4 to our injected classes."""

    def test_get_classes_returns_injected_module(self, applied_patch):
        from mlx_lm.utils import _get_classes

        model_class, args_class = _get_classes({"model_type": "deepseek_v4"})
        assert model_class.__module__ == "mlx_lm.models.deepseek_v4"
        assert args_class.__module__ == "mlx_lm.models.deepseek_v4"
        assert model_class.__name__ == "Model"
        assert args_class.__name__ == "ModelArgs"

    def test_get_classes_returns_injected_module_for_mtp_variant(self, applied_patch):
        from mlx_lm.utils import _get_classes

        model_class, args_class = _get_classes({"model_type": "deepseek_v4_mtp"})
        assert model_class.__module__ == "mlx_lm.models.deepseek_v4"
        assert args_class.__module__ == "mlx_lm.models.deepseek_v4"


class TestPatchedLoadModelTrustRemoteCode:
    """DeepSeek's patched load_model must mirror mlx-lm's custom-code gate."""

    def test_signature_accepts_trust_remote_code(self, applied_patch):
        from mlx_lm.utils import load_model

        assert "trust_remote_code" in inspect.signature(load_model).parameters

    def test_model_file_requires_trust_remote_code(self, tmp_path, applied_patch):
        config_path = tmp_path / "config.json"
        config_path.write_text(
            '{"model_type": "custom", "model_file": "custom_arch.py"}'
        )
        (tmp_path / "custom_arch.py").write_text(
            "\n".join(
                [
                    "from pathlib import Path",
                    "import mlx.nn as nn",
                    "Path(__file__).with_name('executed.txt').write_text('yes')",
                    "",
                    "class ModelArgs:",
                    "    @classmethod",
                    "    def from_dict(cls, config):",
                    "        return cls()",
                    "",
                    "class Model(nn.Module):",
                    "    def __init__(self, args):",
                    "        super().__init__()",
                ]
            )
        )

        from mlx_lm.utils import load_model

        with pytest.raises(ValueError, match="trust_remote_code=True"):
            load_model(tmp_path, strict=False, lazy=True)

        assert not (tmp_path / "executed.txt").exists()

        load_model(
            tmp_path,
            strict=False,
            lazy=True,
            trust_remote_code=True,
        )
        assert (tmp_path / "executed.txt").read_text() == "yes"


class TestCacheHandlerRegistration:
    """omlx CacheTypeRegistry resolves the new cache types to their handlers."""

    def test_pooling_cache_resolves_to_handler(self, applied_patch):
        from omlx.cache.type_registry import CacheTypeRegistry

        handler = CacheTypeRegistry.get_handler_by_class_name("PoolingCache")
        assert type(handler).__name__ == "PoolingCacheHandler"

    def test_batch_pooling_cache_resolves_to_handler(self, applied_patch):
        from omlx.cache.type_registry import CacheTypeRegistry

        handler = CacheTypeRegistry.get_handler_by_class_name("BatchPoolingCache")
        assert type(handler).__name__ == "BatchPoolingCacheHandler"

    def test_pooling_cache_not_block_sliceable(self, applied_patch):
        from omlx.cache.type_registry import CacheTypeRegistry

        handler = CacheTypeRegistry.get_handler_by_class_name("PoolingCache")
        assert handler.supports_block_slicing is False

    def test_batch_pooling_cache_not_block_sliceable(self, applied_patch):
        from omlx.cache.type_registry import CacheTypeRegistry

        handler = CacheTypeRegistry.get_handler_by_class_name("BatchPoolingCache")
        assert handler.supports_block_slicing is False

    def test_detect_cache_type_pooling(self, applied_patch):
        from mlx_lm.models.cache import PoolingCache

        from omlx.cache.type_handlers import CacheType
        from omlx.cache.type_registry import CacheTypeRegistry

        cache = PoolingCache(ratio=4)
        assert CacheTypeRegistry.detect_cache_type(cache) == CacheType.POOLING_CACHE


class TestPoolingCacheStateRoundTrip:
    """Handler extract_state → reconstruct_cache must preserve the pool tensor."""

    def test_round_trip_with_pooled_tensor(self, applied_patch):
        import mlx.core as mx
        from mlx_lm.models.cache import PoolingCache

        from omlx.cache.type_registry import CacheTypeRegistry

        # Build a PoolingCache with a known pool.
        ratio = 4
        cache = PoolingCache(ratio=ratio)
        # Simulate update_and_fetch having stuffed the pool with 8
        # compressed tokens of dim 32.
        pooled = mx.arange(1 * 8 * 32, dtype=mx.float32).reshape(1, 8, 32)
        cache.pooled = pooled

        handler = CacheTypeRegistry.get_handler_by_class_name("PoolingCache")
        state = handler.extract_state(cache)
        assert state["pooled"] is not None
        assert state["pooled"].shape == (1, 8, 32)

        restored = handler.reconstruct_cache(state, meta_state=ratio)
        assert restored is not None
        assert restored.ratio == ratio
        assert restored.pooled.shape == (1, 8, 32)
        # Verify content matches.
        diff = mx.max(mx.abs(restored.pooled - pooled)).item()
        assert diff == 0.0

    def test_round_trip_empty_cache(self, applied_patch):
        from mlx_lm.models.cache import PoolingCache

        from omlx.cache.type_registry import CacheTypeRegistry

        cache = PoolingCache(ratio=8)
        handler = CacheTypeRegistry.get_handler_by_class_name("PoolingCache")
        state = handler.extract_state(cache)
        assert state["pooled"] is None
        assert state["buf_kv"] is None

        restored = handler.reconstruct_cache(state, meta_state=8)
        assert restored is not None
        assert restored.empty()
        assert restored.ratio == 8

    def test_seq_len_from_state(self, applied_patch):
        import mlx.core as mx
        from mlx_lm.models.cache import PoolingCache

        from omlx.cache.type_registry import CacheTypeRegistry

        cache = PoolingCache(ratio=4)
        cache.pooled = mx.zeros((1, 12, 16), dtype=mx.float32)
        handler = CacheTypeRegistry.get_handler_by_class_name("PoolingCache")
        state = handler.extract_state(cache)
        assert handler.get_seq_len(state) == 12


class TestCacheMaterialization:
    """DeepSeek-V4 cache arrays are materialized after forward updates."""

    def test_helper_collects_plain_and_cachelist_leaf_arrays(
        self, applied_patch, monkeypatch
    ):
        import mlx.core as mx
        from mlx_lm.models.cache import CacheList

        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]

        class Leaf:
            def __init__(self, arr):
                self.arr = arr
                self.none_value = None
                self.scalar = 7

        leaf_a = Leaf(mx.array([1], dtype=mx.int32))
        leaf_b = Leaf(mx.array([2], dtype=mx.int32))
        leaf_c = Leaf(mx.array([3], dtype=mx.int32))
        calls = []

        def fake_eval(*arrays):
            calls.append(arrays)

        monkeypatch.setattr(dsv4.mx, "eval", fake_eval)

        dsv4._materialize_cache_arrays([CacheList(leaf_a, leaf_b), leaf_c, None])

        assert len(calls) == 1
        assert calls[0] == (leaf_a.arr, leaf_b.arr, leaf_c.arr)

    def test_model_call_materializes_cache_after_layer_loop(self, applied_patch):
        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]
        source = inspect.getsource(dsv4.DeepseekV4Model.__call__)

        loop_pos = source.index("for layer, layer_cache in zip")
        materialize_pos = source.index("_materialize_cache_arrays(cache)")
        pipeline_send_pos = source.index("if pipeline_rank != 0")

        assert loop_pos < materialize_pos < pipeline_send_pos


class TestDeepseekV4SwitchGLU:
    """DeepSeek-V4 SwitchGLU execution guards."""

    def test_shared_expert_uses_configured_swiglu_limit(self, applied_patch):
        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]

        config = dsv4.ModelArgs(
            vocab_size=16,
            hidden_size=8,
            intermediate_size=16,
            moe_intermediate_size=4,
            num_hidden_layers=1,
            num_attention_heads=2,
            num_key_value_heads=1,
            n_shared_experts=1,
            n_routed_experts=2,
            num_experts_per_tok=1,
            num_hash_layers=0,
            q_lora_rank=0,
            qk_rope_head_dim=4,
            head_dim=4,
            o_lora_rank=0,
            index_n_heads=2,
            index_head_dim=4,
            index_topk=2,
            swiglu_limit=10.0,
        )

        moe = dsv4.DeepseekV4MoE(config, layer_idx=0)

        assert moe.switch_mlp.activation.limit == config.swiglu_limit
        assert moe.shared_experts.swiglu_limit == config.swiglu_limit

    def test_skips_fused_weighted_sum_for_cache_stability(
        self, applied_patch, monkeypatch
    ):
        mx = pytest.importorskip("mlx.core")
        from omlx.patches.deepseek_v4 import switch_layers

        monkeypatch.setattr(
            switch_layers.glm_fast,
            "has_symbol",
            lambda name: name == "glm_moe_weighted_sum",
        )

        def fail_weighted_sum(*args, **kwargs):
            raise AssertionError("DeepSeek V4 must not use fused weighted sum")

        monkeypatch.setattr(
            switch_layers.glm_fast,
            "glm_moe_weighted_sum",
            fail_weighted_sum,
            raising=False,
        )

        mx.random.seed(11)
        layer = switch_layers.SwitchGLU(
            input_dims=16,
            hidden_dims=32,
            num_experts=4,
            bias=False,
        )
        x = mx.random.normal((1, 8, 16), dtype=mx.float32)
        indices = mx.array(
            [
                [
                    [0, 1, 2, 3, 0, 1, 2, 3],
                    [1, 2, 3, 0, 1, 2, 3, 0],
                    [2, 3, 0, 1, 2, 3, 0, 1],
                    [3, 0, 1, 2, 3, 0, 1, 2],
                    [0, 2, 1, 3, 0, 2, 1, 3],
                    [1, 3, 2, 0, 1, 3, 2, 0],
                    [2, 0, 3, 1, 2, 0, 3, 1],
                    [3, 1, 0, 2, 3, 1, 0, 2],
                ]
            ],
            dtype=mx.int32,
        )
        scores = mx.softmax(
            mx.random.normal((1, 8, 8), dtype=mx.float32),
            axis=-1,
        )

        y = layer(x, indices, scores=scores)
        mx.eval(y)

        assert y.shape == (1, 8, 8, 16)


class TestPreLoadDispatch:
    """maybe_apply_pre_load_patches gates correctly on config.json model_type."""

    def test_no_dispatch_for_other_model_type(self, tmp_path):
        # Create a fake model dir with a non-deepseek config.
        config_path = tmp_path / "config.json"
        config_path.write_text('{"model_type": "llama"}')

        from omlx.utils.model_loading import maybe_apply_pre_load_patches

        # Should be a no-op (no exception). We can't easily assert that
        # apply_deepseek_v4_patch was NOT called because earlier tests
        # may have applied it already. Just verify no crash.
        maybe_apply_pre_load_patches(str(tmp_path))

    def test_no_dispatch_for_missing_config(self, tmp_path):
        # No config.json present.
        from omlx.utils.model_loading import maybe_apply_pre_load_patches

        maybe_apply_pre_load_patches(str(tmp_path))

    def test_dispatch_for_deepseek_v4(self, tmp_path):
        config_path = tmp_path / "config.json"
        config_path.write_text('{"model_type": "deepseek_v4"}')

        from omlx.patches.deepseek_v4 import is_applied
        from omlx.utils.model_loading import maybe_apply_pre_load_patches

        maybe_apply_pre_load_patches(str(tmp_path))
        # Patch must be applied after this dispatch (or already applied).
        assert is_applied() is True

    def test_dispatch_for_deepseek_v4_mtp_variant(self, tmp_path):
        config_path = tmp_path / "config.json"
        config_path.write_text('{"model_type": "deepseek_v4_mtp"}')

        from omlx.patches.deepseek_v4 import is_applied
        from omlx.utils.model_loading import maybe_apply_pre_load_patches

        maybe_apply_pre_load_patches(str(tmp_path))
        assert is_applied() is True


class TestMakeQuantizationConfigMtp:
    """make_quantization_config must cover the MTP fusion projections.

    Without explicit entries, mtp.<i>.e_proj / mtp.<i>.h_proj fall through
    to the affine default, whose QuantizedLinear expects a .biases tensor
    the fp8 checkpoint doesn't ship, and strict load fails."""

    def test_mtp_projections_get_mxfp8(self, applied_patch):
        import mlx.nn as nn

        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]

        class _MTPStub(nn.Module):
            def __init__(self):
                super().__init__()
                self.e_proj = nn.Linear(8, 8, bias=False)
                self.h_proj = nn.Linear(8, 8, bias=False)

        class _ModelStub(nn.Module):
            def __init__(self):
                super().__init__()
                self.mtp = [_MTPStub()]
                self.lm_head = nn.Linear(8, 8, bias=False)

        qcfg = dsv4.make_quantization_config(_ModelStub())
        mxfp8 = {"group_size": 32, "bits": 8, "mode": "mxfp8"}
        assert qcfg["mtp.0.e_proj"] == mxfp8
        assert qcfg["mtp.0.h_proj"] == mxfp8
        # Non-MTP paths keep the affine default (no per-path entry).
        assert "lm_head" not in qcfg

    def test_no_mtp_no_entries(self, applied_patch):
        import mlx.nn as nn

        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]

        class _ModelStub(nn.Module):
            def __init__(self):
                super().__init__()
                self.lm_head = nn.Linear(8, 8, bias=False)

        qcfg = dsv4.make_quantization_config(_ModelStub())
        assert not any(k.startswith("mtp.") for k in qcfg)


class TestDeepSeekV4SanitizeAffineSwitchMLP:
    """Sanitize should enable the FP16 affine routed-MoE fast path."""

    def test_affine_switch_mlp_scale_bias_cast_to_fp16(self, applied_patch):
        mx = pytest.importorskip("mlx.core")

        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]
        fake_model = SimpleNamespace(
            args=SimpleNamespace(
                num_hidden_layers=1,
                n_routed_experts=2,
                o_groups=1,
                o_lora_rank=1,
            )
        )
        weights = {
            "model.layers.0.ffn.switch_mlp.up_proj.weight": mx.zeros(
                (2, 4, 2), dtype=mx.uint32
            ),
            "model.layers.0.ffn.switch_mlp.up_proj.scales": mx.zeros(
                (2, 4, 1), dtype=mx.bfloat16
            ),
            "model.layers.0.ffn.switch_mlp.up_proj.biases": mx.zeros(
                (2, 4, 1), dtype=mx.bfloat16
            ),
            "model.layers.0.ffn.switch_mlp.down_proj.weight": mx.zeros(
                (2, 4, 2), dtype=mx.uint32
            ),
            "model.layers.0.ffn.switch_mlp.down_proj.scales": mx.zeros(
                (2, 4, 1), dtype=mx.bfloat16
            ),
            "model.layers.0.ffn.switch_mlp.down_proj.biases": mx.zeros(
                (2, 4, 1), dtype=mx.bfloat16
            ),
            "model.layers.0.ffn.shared_experts.up_proj.scales": mx.zeros(
                (4, 1), dtype=mx.bfloat16
            ),
        }

        out = dsv4.Model.sanitize(fake_model, dict(weights))

        assert out["model.layers.0.ffn.switch_mlp.up_proj.scales"].dtype == mx.float16
        assert out["model.layers.0.ffn.switch_mlp.up_proj.biases"].dtype == mx.float16
        assert out["model.layers.0.ffn.switch_mlp.down_proj.scales"].dtype == mx.float16
        assert out["model.layers.0.ffn.switch_mlp.down_proj.biases"].dtype == mx.float16
        assert (
            out["model.layers.0.ffn.shared_experts.up_proj.scales"].dtype == mx.bfloat16
        )


class TestDeepSeekV4SanitizeHcAliases:
    """Sanitize accepts both upstream HC key spellings for V4 checkpoints."""

    @staticmethod
    def _fake_model():
        return SimpleNamespace(
            args=SimpleNamespace(
                num_hidden_layers=1,
                n_routed_experts=0,
                o_groups=1,
                o_lora_rank=1,
            )
        )

    def test_dotted_hc_aliases_remap_to_model_modules(self, applied_patch):
        mx = pytest.importorskip("mlx.core")

        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]
        weights = {
            "model.layers.0.hc_attn.base": mx.zeros((1,), dtype=mx.float32),
            "model.layers.0.hc_attn.fn": mx.zeros((1, 1), dtype=mx.float32),
            "model.layers.0.hc_attn.scale": mx.zeros((3,), dtype=mx.float32),
            "model.layers.0.hc_ffn.base": mx.zeros((1,), dtype=mx.float32),
            "model.layers.0.hc_ffn.fn": mx.zeros((1, 1), dtype=mx.float32),
            "model.layers.0.hc_ffn.scale": mx.zeros((3,), dtype=mx.float32),
        }

        out = dsv4.Model.sanitize(self._fake_model(), dict(weights))

        assert "model.layers.0.attn_hc.base" in out
        assert "model.layers.0.attn_hc.fn" in out
        assert "model.layers.0.attn_hc.scale" in out
        assert "model.layers.0.ffn_hc.base" in out
        assert "model.layers.0.ffn_hc.fn" in out
        assert "model.layers.0.ffn_hc.scale" in out
        assert not any(".hc_attn." in key or ".hc_ffn." in key for key in out)

    def test_dotted_hc_alias_does_not_override_canonical_key(self, applied_patch):
        mx = pytest.importorskip("mlx.core")

        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]
        weights = {
            "model.layers.0.hc_attn.base": mx.zeros((1,), dtype=mx.float32),
            "model.layers.0.attn_hc.base": mx.zeros((2,), dtype=mx.float32),
        }

        out = dsv4.Model.sanitize(self._fake_model(), dict(weights))

        assert out["model.layers.0.attn_hc.base"].shape == (2,)
        assert "model.layers.0.hc_attn.base" not in out


class TestMtpSanitizeWoAReshape:
    """The MTP patch sanitize must reshape mtp.<i>.block.attn.wo_a from the
    2D nn.Linear layout to the 3D MultiLinear layout, like the backbone."""

    @pytest.fixture()
    def patched_sanitize(self, applied_patch):
        import omlx.patches.mlx_lm_mtp.deepseek_v4_model as mtp_dsv4

        mtp_dsv4.apply()
        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]
        return dsv4.Model.sanitize

    @staticmethod
    def _fake_model(with_mtp=True):
        class _Args:
            num_hidden_layers = 1
            num_nextn_predict_layers = 1
            o_groups = 2
            o_lora_rank = 4
            n_routed_experts = 2

        class _Fake:
            args = _Args()

        fake = _Fake()
        if with_mtp:
            fake.mtp = [object()]
        return fake

    def test_mtp_wo_a_2d_reshaped_to_3d(self, patched_sanitize):
        import mlx.core as mx

        weights = {
            "mtp.0.attn.wo_a.weight": mx.zeros((8, 16), dtype=mx.bfloat16),
        }
        out = patched_sanitize(self._fake_model(), weights)
        assert out["mtp.0.block.attn.wo_a.weight"].shape == (2, 4, 16)

    def test_mtp_wo_a_3d_unchanged(self, patched_sanitize):
        import mlx.core as mx

        weights = {
            "mtp.0.block.attn.wo_a.weight": mx.zeros((2, 4, 16), dtype=mx.bfloat16),
        }
        out = patched_sanitize(self._fake_model(), weights)
        assert out["mtp.0.block.attn.wo_a.weight"].shape == (2, 4, 16)

    def test_mtp_dotted_hc_alias_nested_under_block(self, patched_sanitize):
        import mlx.core as mx

        weights = {
            "mtp.0.hc_attn.base": mx.zeros((1,), dtype=mx.float32),
            "mtp.0.hc_ffn.scale": mx.zeros((3,), dtype=mx.float32),
        }

        out = patched_sanitize(self._fake_model(), weights)

        assert "mtp.0.block.attn_hc.base" in out
        assert "mtp.0.block.ffn_hc.scale" in out
        assert "mtp.0.hc_attn.base" not in out
        assert "mtp.0.hc_ffn.scale" not in out


class TestMtpBackboneInterface:
    """The patched DSv4 Model.__call__ must accept the full patched-backbone
    interface — batch_generator._call_backbone passes n_confirmed=1 during
    MTP verify cycles (crashed with TypeError before the fix)."""

    def test_call_accepts_n_confirmed(self, applied_patch):
        import omlx.patches.mlx_lm_mtp.deepseek_v4_model as mtp_dsv4

        mtp_dsv4.apply()
        dsv4 = sys.modules["mlx_lm.models.deepseek_v4"]
        sig = inspect.signature(dsv4.Model.__call__)
        assert "n_confirmed" in sig.parameters
        assert sig.parameters["n_confirmed"].default == 0
        assert "return_hidden" in sig.parameters


class TestPoolingCacheTrimRollback:
    """trim(1) must exactly undo the last (draft) token of an MTP verify
    update, including the pool-boundary case where the draft completed a
    compression window. Equivalence is checked behaviorally: a trimmed
    cache must evolve identically to a reference cache that never saw the
    rejected token."""

    @staticmethod
    def _push(cache, tokens, offset):
        """Feed raw per-token rows through the PoolingCache contract,
        compressing completed windows with a deterministic stand-in
        (mean over the window) like Compressor does."""
        import mlx.core as mx

        kv = tokens
        gate = tokens * 0.5
        r_kv, _r_gate, _ = cache.accumulate_windows(kv, gate, offset)
        if r_kv.size == 0:
            rows = mx.zeros((kv.shape[0], 0, kv.shape[-1]), dtype=kv.dtype)
        else:
            rows = mx.unflatten(r_kv, 1, (-1, cache.ratio)).mean(axis=2)
        return cache.update_and_fetch(rows)

    @staticmethod
    def _tok(values):
        import mlx.core as mx

        arr = mx.array(values, dtype=mx.float32)
        return mx.broadcast_to(arr[None, :, None], (1, len(values), 8))

    def _equivalence(self, cache_cls, prefix, verify, post, applied):
        """Drive cache through prefix + 2-token verify, trim the draft,
        push `post`; compare against a reference that never saw the draft."""
        import mlx.core as mx

        ratio = 4
        if cache_cls.__name__ == "BatchPoolingCache":
            cache = cache_cls(ratio, [0])
            ref = cache_cls(ratio, [0])
        else:
            cache = cache_cls(ratio)
            ref = cache_cls(ratio)

        pos = 0
        for chunk in prefix:
            self._push(cache, self._tok(chunk), pos)
            self._push(ref, self._tok(chunk), pos)
            pos += len(chunk)

        # Verify forward: [confirmed, draft] on cache; confirmed only on ref.
        self._push(cache, self._tok(verify), pos)
        assert cache.is_trimmable()
        assert cache.trim(1) == 1
        self._push(ref, self._tok(verify[:1]), pos)
        pos += 1

        out = self._push(cache, self._tok(post), pos)
        ref_out = self._push(ref, self._tok(post), pos)

        if out is None or getattr(out, "size", 0) == 0:
            assert ref_out is None or getattr(ref_out, "size", 0) == 0
        else:
            pl = getattr(cache, "_pool_lengths", None)
            n = pl[0] if pl is not None else out.shape[1]
            ref_n = ref._pool_lengths[0] if pl is not None else ref_out.shape[1]
            assert n == ref_n
            assert mx.allclose(out[:, :n], ref_out[:, :n]).item()
        assert (
            cache.remainder
            if isinstance(cache.remainder, int)
            else list(cache.remainder)
        ) == (ref.remainder if isinstance(ref.remainder, int) else list(ref.remainder))

    def test_easy_case_draft_in_buffer(self, applied_patch):
        from mlx_lm.models.cache import PoolingCache

        # After verify: remainder = (1 + 2) % 4 = 3 >= 1 -> buffer trim.
        self._equivalence(PoolingCache, [[1.0]], [2.0, 3.0], [4.0], applied_patch)

    def test_boundary_case_draft_completed_window(self, applied_patch):
        from mlx_lm.models.cache import PoolingCache

        # remainder before verify = 2; verify adds 2 -> window completes on
        # the draft token -> undo log path (drop pooled row, replay
        # confirmed into the buffer).
        self._equivalence(
            PoolingCache, [[1.0, 2.0]], [3.0, 4.0], [5.0, 6.0, 7.0], applied_patch
        )

    def test_boundary_case_with_existing_pool(self, applied_patch):
        from mlx_lm.models.cache import PoolingCache

        # One full window already pooled, then the boundary case again.
        self._equivalence(
            PoolingCache,
            [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0]],
            [7.0, 8.0],
            [9.0, 10.0, 11.0],
            applied_patch,
        )

    def test_batch_easy_case(self, applied_patch):
        from mlx_lm.models.cache import BatchPoolingCache

        self._equivalence(BatchPoolingCache, [[1.0]], [2.0, 3.0], [4.0], applied_patch)

    def test_batch_boundary_case(self, applied_patch):
        from mlx_lm.models.cache import BatchPoolingCache

        self._equivalence(
            BatchPoolingCache,
            [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0]],
            [7.0, 8.0],
            [9.0, 10.0, 11.0],
            applied_patch,
        )

    def test_untrimmable_when_no_undo_after_prompt(self, applied_patch):
        """Prompt-sized updates (L > 8) don't stash an undo log; a trim at
        a pool boundary right after one must report not-trimmable instead
        of corrupting state. (Updates up to L == 8 keep an undo so depth-k
        MTP verify windows can roll back.)"""
        from mlx_lm.models.cache import PoolingCache

        cache = PoolingCache(4)
        self._push(
            cache,
            self._tok([float(v) for v in range(1, 13)]),  # L = 12 > 8
            0,
        )
        assert cache.remainder == 0
        assert cache.pooled is not None
        assert not cache.is_trimmable()
        assert cache.trim(1) == 0

    def test_verify_sized_update_keeps_undo(self, applied_patch):
        """MTP verify windows (2 < L <= 8) stash an undo log: a trim right
        after one rolls back across the pool boundary instead of failing."""
        from mlx_lm.models.cache import PoolingCache

        cache = PoolingCache(4)
        self._push(cache, self._tok([1.0, 2.0, 3.0, 4.0]), 0)
        assert cache.remainder == 0
        assert cache.pooled is not None
        assert cache.is_trimmable()
        assert cache.trim(1) == 1
        # The completed window is undone: its 3 surviving tokens are back
        # in the remainder buffer and no pooled row remains visible.
        assert cache.remainder == 3
        assert cache.size() == 0


class TestNaxMoEStockRouting:
    """NAX GPUs route prefill-sized MoE gemms to stock mx.gather_qmm."""

    @pytest.fixture(autouse=True)
    def _nax_off_by_default(self, monkeypatch):
        from omlx.patches.deepseek_v4 import switch_layers as sl

        # Pin detection off so the block-kernel tests behave identically on
        # M5-family machines; each test overrides what it needs.
        monkeypatch.setattr(sl, "is_nax_available", lambda: False)
        monkeypatch.setattr(sl, "_NAX_STOCK_MODE", "")
        yield

    def test_prefers_stock_for_prefill_route_counts_only(self, monkeypatch):
        from omlx.patches.deepseek_v4 import switch_layers as sl

        monkeypatch.setattr(sl, "is_nax_available", lambda: True)
        assert not sl._nax_prefers_stock(8)
        assert not sl._nax_prefers_stock(sl._NAX_STOCK_MIN_ROUTES - 1)
        assert sl._nax_prefers_stock(sl._NAX_STOCK_MIN_ROUTES)
        assert sl._nax_prefers_stock(1 << 20)

    def test_no_stock_routing_without_nax(self, monkeypatch):
        from omlx.patches.deepseek_v4 import switch_layers as sl

        assert not sl._nax_prefers_stock(1 << 20)

    def test_env_kill_switch_keeps_block_kernels(self, monkeypatch):
        from omlx.patches.deepseek_v4 import switch_layers as sl

        monkeypatch.setattr(sl, "is_nax_available", lambda: True)
        monkeypatch.setattr(sl, "_NAX_STOCK_MODE", "0")
        assert not sl._nax_prefers_stock(1 << 20)

    def test_env_force_routes_everything(self, monkeypatch):
        from omlx.patches.deepseek_v4 import switch_layers as sl

        monkeypatch.setattr(sl, "is_nax_available", lambda: True)
        monkeypatch.setattr(sl, "_NAX_STOCK_MODE", "1")
        assert sl._nax_prefers_stock(1)

    def test_native_block_kind_short_circuits_on_nax_prefill(self, monkeypatch):
        import mlx.core as mx

        from omlx.patches.deepseek_v4 import switch_layers as sl

        linear = sl.QuantizedSwitchLinear(
            64, 64, num_experts=2, bias=False, group_size=64, bits=4
        )
        monkeypatch.setattr(sl, "_nax_prefers_stock", lambda n: n >= 1024)
        prefill_x = mx.zeros((2048, 1, 64), dtype=mx.bfloat16)
        assert linear._native_block_kind(prefill_x, True) is None
        # Decode-sized calls fall through to the regular block-kernel gates:
        # the NAX gate must not change what they resolve to.
        decode_x = mx.zeros((8, 1, 64), dtype=mx.bfloat16)
        gated = linear._native_block_kind(decode_x, True)
        monkeypatch.setattr(sl, "_nax_prefers_stock", lambda n: False)
        assert gated == linear._native_block_kind(decode_x, True)