# SPDX-License-Identifier: Apache-2.0
"""Tests for omlx.speculative.vlm_mtp.

Phase 2A: covers drafter validation, lazy bind, and wrapper-level dispatch
to mlx-vlm's ``_mtp_rounds`` / ``_mtp_rounds_batch``. The actual mlx-vlm
helpers are mocked so this suite stays fast and does not touch model
weights.
"""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import mlx.core as mx
import pytest

from omlx.speculative import vlm_mtp


def _fake_drafter_model(model_type: str = "gemma4_assistant") -> MagicMock:
    """Build a stand-in for Gemma4AssistantDraftModel that satisfies the
    minimum API used by VLMMTPDrafter."""
    drafter = MagicMock()
    drafter.config = MagicMock(model_type=model_type)
    return drafter


def test_load_vlm_mtp_drafter_happy_path():
    """Valid gemma4_assistant artifact returns a populated VLMMTPDrafter."""
    fake_model = _fake_drafter_model("gemma4_assistant")
    with patch.object(vlm_mtp, "_vlm_load_drafter", return_value=(fake_model, "mtp")):
        drafter = vlm_mtp.load_vlm_mtp_drafter("/path/to/drafter")
    assert isinstance(drafter, vlm_mtp.VLMMTPDrafter)
    assert drafter.draft_kind == "mtp"
    assert drafter.source_path == "/path/to/drafter"
    assert drafter.model is fake_model


def test_load_vlm_mtp_drafter_accepts_unified_assistant():
    """Valid gemma4_unified_assistant artifact is accepted."""
    fake_model = _fake_drafter_model("gemma4_unified_assistant")
    with patch.object(vlm_mtp, "_vlm_load_drafter", return_value=(fake_model, "mtp")):
        drafter = vlm_mtp.load_vlm_mtp_drafter("/path/to/drafter")
    assert isinstance(drafter, vlm_mtp.VLMMTPDrafter)
    assert drafter.model is fake_model


def test_load_vlm_mtp_drafter_rejects_dflash_kind():
    """A drafter that resolves to non-mtp kind is rejected (None + warn)."""
    fake_model = _fake_drafter_model("qwen3_dflash")
    with patch.object(
        vlm_mtp, "_vlm_load_drafter", return_value=(fake_model, "dflash")
    ):
        result = vlm_mtp.load_vlm_mtp_drafter("/path/to/drafter")
    assert result is None


def test_load_vlm_mtp_drafter_accepts_qwen3_5_mtp():
    """qwen3_5_mtp model_type with kind='mtp' is accepted."""
    fake_model = _fake_drafter_model("qwen3_5_mtp")
    with patch.object(vlm_mtp, "_vlm_load_drafter", return_value=(fake_model, "mtp")):
        drafter = vlm_mtp.load_vlm_mtp_drafter("/path/to/qwen-mtp")
    assert isinstance(drafter, vlm_mtp.VLMMTPDrafter)
    assert drafter.draft_kind == "mtp"
    assert drafter.model is fake_model


def test_load_vlm_mtp_drafter_swallows_load_exception():
    """Load failures are logged and converted to None — never raised."""
    with patch.object(
        vlm_mtp,
        "_vlm_load_drafter",
        side_effect=RuntimeError("HF repo not found"),
    ):
        result = vlm_mtp.load_vlm_mtp_drafter("not-a-real-drafter")
    assert result is None


def test_run_vlm_mtp_decode_single_request_dispatches_to_mtp_rounds():
    """Single-int first_bonus routes to ``_mtp_rounds``, yields first_bonus
    then any tokens that the round loop emits."""
    fake_model = _fake_drafter_model("gemma4_assistant")
    drafter = vlm_mtp.VLMMTPDrafter(fake_model, "mtp", "/p")
    target = MagicMock()
    sampler = MagicMock()

    yielded = [(11, None), (22, None), (33, None)]
    with (
        patch.object(vlm_mtp, "_mtp_rounds", return_value=iter(yielded)) as m_single,
        patch.object(vlm_mtp, "_mtp_rounds_batch") as m_batch,
        patch.object(vlm_mtp, "_buffer_mtp_target_cache") as m_buffer,
    ):
        prompt_tokens = mx.array([[5, 6, 7]], dtype=mx.int32)
        out = list(
            vlm_mtp.run_vlm_mtp_decode(
                target_language_model=target,
                drafter=drafter,
                prompt_cache=[],
                hidden=mx.zeros((1, 1, 8)),
                shared_kv_states={},
                first_bonus=7,
                max_tokens=4,
                sampler=sampler,
                prompt_tokens=prompt_tokens,
            )
        )

    # first_bonus 7 is yielded by the wrapper before _mtp_rounds takes over
    assert out == [7, 11, 22, 33]
    m_single.assert_called_once()
    m_batch.assert_not_called()
    m_buffer.assert_called_once()
    buffer_args = m_buffer.call_args.args
    assert buffer_args[0] == []
    assert getattr(buffer_args[1], "_drafter", buffer_args[1]) is fake_model
    assert buffer_args[2] is None
    # first_bonus int forwarded as int
    kwargs = m_single.call_args.kwargs
    assert kwargs["first_bonus"] == 7
    assert kwargs["max_tokens"] == 4
    assert kwargs["prompt_tokens"] is prompt_tokens


def test_run_vlm_mtp_decode_batch_dispatches_to_mtp_rounds_batch():
    """Multi-row mx.array first_bonus routes to ``_mtp_rounds_batch``,
    emits first_bonus row then the round-loop rows."""
    fake_model = _fake_drafter_model("gemma4_assistant")
    drafter = vlm_mtp.VLMMTPDrafter(fake_model, "mtp", "/p")
    target = MagicMock()
    sampler = MagicMock()

    first_bonus = mx.array([1, 2, 3])  # B=3
    yielded = [([1, None, 3], None), ([None, None, None], None)]
    with (
        patch.object(
            vlm_mtp, "_mtp_rounds_batch", return_value=iter(yielded)
        ) as m_batch,
        patch.object(vlm_mtp, "_mtp_rounds") as m_single,
        patch.object(vlm_mtp, "_buffer_mtp_target_cache") as m_buffer,
    ):
        out = list(
            vlm_mtp.run_vlm_mtp_decode(
                target_language_model=target,
                drafter=drafter,
                prompt_cache=[],
                hidden=mx.zeros((3, 1, 8)),
                shared_kv_states={},
                first_bonus=first_bonus,
                max_tokens=4,
                sampler=sampler,
                eos_token_ids={2, 5},
            )
        )

    # First yielded row is the first_bonus row (one int per request).
    assert out == [[1, 2, 3], [1, None, 3], [None, None, None]]
    m_batch.assert_called_once()
    m_single.assert_not_called()
    m_buffer.assert_not_called()
    kwargs = m_batch.call_args.kwargs
    # EOS forwarded as a fresh set (function does its own copy)
    assert kwargs["eos_token_ids"] == {2, 5}


def test_run_vlm_mtp_decode_single_scalar_array_unwraps_to_int():
    """B=1 mx.array first_bonus is treated as single-request and unwrapped."""
    fake_model = _fake_drafter_model("gemma4_assistant")
    drafter = vlm_mtp.VLMMTPDrafter(fake_model, "mtp", "/p")
    target = MagicMock()
    sampler = MagicMock()

    first_bonus = mx.array([42])  # B=1 should not take the batch branch
    with (
        patch.object(vlm_mtp, "_mtp_rounds", return_value=iter([])) as m_single,
        patch.object(vlm_mtp, "_mtp_rounds_batch") as m_batch,
    ):
        out = list(
            vlm_mtp.run_vlm_mtp_decode(
                target_language_model=target,
                drafter=drafter,
                prompt_cache=[],
                hidden=mx.zeros((1, 1, 8)),
                shared_kv_states={},
                first_bonus=first_bonus,
                max_tokens=4,
                sampler=sampler,
            )
        )

    # _mtp_rounds yields nothing here, so only the wrapper's first_bonus
    # emit makes it into the stream.
    assert out == [42]
    m_single.assert_called_once()
    m_batch.assert_not_called()
    assert m_single.call_args.kwargs["first_bonus"] == 42


@pytest.mark.parametrize(
    "vlm_mtp_kw, other_kw",
    [
        ("dflash_enabled", "dflash_enabled"),
        ("specprefill_enabled", "specprefill_enabled"),
        ("mtp_enabled", "mtp_enabled"),
        ("turboquant_kv_enabled", "turboquant_kv_enabled"),
    ],
)
def test_model_settings_vlm_mtp_mutex(vlm_mtp_kw, other_kw):
    """ModelSettings.__post_init__ raises when vlm_mtp_enabled overlaps
    with any other speculative / cache-mutating toggle."""
    from omlx.model_settings import ModelSettings

    with pytest.raises(ValueError, match="vlm_mtp_enabled"):
        ModelSettings(vlm_mtp_enabled=True, **{other_kw: True})


# ---------------------------------------------------------------------------
# MoE config patch tests
# ---------------------------------------------------------------------------


class TestMoeConfigPatch:
    """Verify that the MoE compat patch in vlm_mtp.py correctly handles
    qwen3_5_moe_text text_config dicts."""

    def test_patch_is_applied_on_import(self):
        """The patch runs at import time; Qwen3_5MTPConfig.__post_init__
        should be the patched version."""
        try:
            from mlx_vlm.speculative.drafters.qwen3_5_mtp.config import (
                Qwen3_5MTPConfig,
            )
        except ImportError:
            pytest.skip("mlx-vlm qwen3_5_mtp drafter not available")

        # The patched __post_init__ is a closure, not the original method.
        # Verify it was replaced by checking it's not the unpatched version.
        src = Qwen3_5MTPConfig.__post_init__
        # The patched version references MoETextConfig in its closure.
        assert src is not None

    def test_moe_text_config_accepted(self):
        """Qwen3_5MTPConfig.from_dict with a MoE text_config does not raise."""
        try:
            from mlx_vlm.speculative.drafters.qwen3_5_mtp.config import (
                Qwen3_5MTPConfig,
            )
        except ImportError:
            pytest.skip("mlx-vlm qwen3_5_mtp drafter not available")

        moe_config = {
            "model_type": "qwen3_5_mtp",
            "text_config": {
                "model_type": "qwen3_5_moe_text",
                "hidden_size": 64,
                "num_hidden_layers": 2,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "num_experts": 8,
                "num_experts_per_tok": 2,
                "shared_expert_intermediate_size": 128,
                "moe_intermediate_size": 128,
                "rms_norm_eps": 1e-6,
                "vocab_size": 256,
                "max_position_embeddings": 128,
                "linear_num_value_heads": 4,
                "linear_num_key_heads": 4,
                "linear_key_head_dim": 16,
                "linear_value_head_dim": 16,
                "linear_conv_kernel_dim": 4,
                "mtp_num_hidden_layers": 1,
            },
        }
        cfg = Qwen3_5MTPConfig.from_dict(moe_config)
        assert cfg.text_config is not None
        assert cfg.text_config.hidden_size == 64
        assert cfg.text_config.num_experts == 8

    def test_dense_text_config_still_works(self):
        """Qwen3_5MTPConfig.from_dict with a dense text_config still works."""
        try:
            from mlx_vlm.speculative.drafters.qwen3_5_mtp.config import (
                Qwen3_5MTPConfig,
            )
        except ImportError:
            pytest.skip("mlx-vlm qwen3_5_mtp drafter not available")

        dense_config = {
            "model_type": "qwen3_5_mtp",
            "text_config": {
                "model_type": "qwen3_5",
                "hidden_size": 64,
                "intermediate_size": 128,
                "num_hidden_layers": 2,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "rms_norm_eps": 1e-6,
                "vocab_size": 256,
                "max_position_embeddings": 128,
                "linear_num_value_heads": 4,
                "linear_num_key_heads": 4,
                "linear_key_head_dim": 16,
                "linear_value_head_dim": 16,
                "linear_conv_kernel_dim": 4,
                "mtp_num_hidden_layers": 1,
            },
        }
        cfg = Qwen3_5MTPConfig.from_dict(dense_config)
        assert cfg.text_config is not None
        assert cfg.text_config.hidden_size == 64


# ---------------------------------------------------------------------------
# dense Qwen3.5 VLM runtime patch tests
# ---------------------------------------------------------------------------


def test_dense_vlm_runtime_return_hidden_uses_language_model_output_contract():
    """Dense Qwen3.5 VLM MTP verify must satisfy mlx-vlm's output contract."""
    from mlx_vlm.models.base import LanguageModelOutput
    from omlx.patches.mlx_vlm_mtp import qwen35_vlm_runtime

    logits = mx.zeros((1, 2, 16))
    hidden = mx.zeros((1, 2, 8))
    gdn_states = [{"state": "mock"}]

    class FakeStockOutput:
        def __init__(self):
            self.logits = logits
            self.hidden_states = [hidden]
            self.gdn_states = gdn_states

    class FakeLanguageModel:
        def __init__(self, args, config=None):
            self.args = args
            self.config = config
            self.model = SimpleNamespace(layers=[object(), object()])
            self.forward_kwargs = None

        def __call__(
            self,
            inputs,
            inputs_embeds=None,
            mask=None,
            cache=None,
            **kwargs,
        ):
            self.forward_kwargs = kwargs
            return FakeStockOutput()

    q35_lang = SimpleNamespace(LanguageModel=FakeLanguageModel)
    qwen35_vlm_runtime._patch_vlm_language_model(q35_lang)

    model = q35_lang.LanguageModel(
        SimpleNamespace(mtp_num_hidden_layers=0, tie_word_embeddings=True),
        config=None,
    )
    out = model(
        mx.array([[1, 2]], dtype=mx.int32),
        cache=[],
        return_hidden=True,
        return_shared_kv=True,
        capture_layer_ids=[99],
    )

    assert isinstance(out, LanguageModelOutput)
    assert out.logits is logits
    assert out.hidden_states == [hidden]
    assert out.hidden_states[-1] is hidden
    assert out.gdn_states is gdn_states
    assert out.shared_kv_states == {}
    assert model.forward_kwargs["capture_layer_ids"] == [1]


def test_moe_vlm_sanitize_unfuses_gate_up_by_midpoint(monkeypatch):
    """The VLM MoE sanitize patch must preserve upstream midpoint slicing."""
    from omlx.patches.mlx_vlm_mtp import qwen35_moe_vlm_model
    from mlx_vlm.models.qwen3_5_moe import qwen3_5_moe

    monkeypatch.setattr(qwen35_moe_vlm_model, "_APPLIED", False)
    if hasattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched"):
        monkeypatch.delattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched")

    assert qwen35_moe_vlm_model.apply() is True

    fake_self = SimpleNamespace(
        config=SimpleNamespace(
            text_config=SimpleNamespace(
                tie_word_embeddings=False,
                num_hidden_layers=1,
                num_experts=0,
            )
        )
    )
    gate_up = mx.arange(2 * 6 * 3).reshape(2, 6, 3)
    weights = {
        "model.language_model.layers.0.mlp.experts.gate_up_proj": gate_up,
        "model.language_model.layers.0.mlp.experts.down_proj": mx.ones((2, 4, 3)),
    }

    result = qwen3_5_moe.Model.sanitize(fake_self, weights)

    gate_key = "language_model.model.layers.0.mlp.switch_mlp.gate_proj.weight"
    up_key = "language_model.model.layers.0.mlp.switch_mlp.up_proj.weight"
    assert bool(mx.all(result[gate_key] == gate_up[:, :3, :]).item())
    assert bool(mx.all(result[up_key] == gate_up[:, 3:, :]).item())


def test_moe_vlm_runtime_sanitize_unfuses_gate_up_by_midpoint():
    """The runtime sanitize wrapper must not reintroduce the old split path."""
    from omlx.patches.mlx_vlm_mtp import qwen35_moe_vlm_runtime

    class FakeModel:
        pass

    fake_outer = SimpleNamespace(Model=FakeModel)
    qwen35_moe_vlm_runtime._patch_vlm_outer_model_sanitize(fake_outer)

    fake_self = SimpleNamespace(
        config=SimpleNamespace(
            text_config=SimpleNamespace(
                tie_word_embeddings=False,
                num_hidden_layers=1,
                num_experts=0,
            )
        )
    )
    gate_up = mx.arange(2 * 6 * 3).reshape(2, 6, 3)
    weights = {
        "model.language_model.layers.0.mlp.experts.gate_up_proj": gate_up,
        "model.language_model.layers.0.mlp.experts.down_proj": mx.ones((2, 4, 3)),
    }

    result = FakeModel.sanitize(fake_self, weights)

    gate_key = "language_model.model.layers.0.mlp.switch_mlp.gate_proj.weight"
    up_key = "language_model.model.layers.0.mlp.switch_mlp.up_proj.weight"
    assert bool(mx.all(result[gate_key] == gate_up[:, :3, :]).item())
    assert bool(mx.all(result[up_key] == gate_up[:, 3:, :]).item())


def _per_expert_vlm_self(num_experts=2, num_hidden_layers=1):
    return SimpleNamespace(
        config=SimpleNamespace(
            text_config=SimpleNamespace(
                tie_word_embeddings=False,
                num_hidden_layers=num_hidden_layers,
                num_experts=num_experts,
            )
        )
    )


def test_moe_vlm_sanitize_stacks_per_expert_backbone(monkeypatch):
    """Ornith / raw Qwen3.5 ship backbone MoE layers as per-expert tensors.
    The model-level sanitize must stack them into switch_mlp form."""
    from omlx.patches.mlx_vlm_mtp import qwen35_moe_vlm_model
    from mlx_vlm.models.qwen3_5_moe import qwen3_5_moe

    monkeypatch.setattr(qwen35_moe_vlm_model, "_APPLIED", False)
    if hasattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched"):
        monkeypatch.delattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched")
    assert qwen35_moe_vlm_model.apply() is True

    pfx_in = "model.language_model.layers.0.mlp"
    weights = {}
    for e in range(2):
        weights[f"{pfx_in}.experts.{e}.gate_proj.weight"] = mx.zeros((8, 4))
        weights[f"{pfx_in}.experts.{e}.up_proj.weight"] = mx.zeros((8, 4))
        weights[f"{pfx_in}.experts.{e}.down_proj.weight"] = mx.zeros((4, 8))

    result = qwen3_5_moe.Model.sanitize(_per_expert_vlm_self(), weights)

    pfx = "language_model.model.layers.0.mlp"
    assert result[f"{pfx}.switch_mlp.gate_proj.weight"].shape == (2, 8, 4)
    assert result[f"{pfx}.switch_mlp.down_proj.weight"].shape == (2, 4, 8)
    assert not any(f"{pfx}.experts." in k for k in result)


def test_moe_vlm_sanitize_stacks_per_expert_backbone_quantized(monkeypatch):
    """A per-expert *quantized* backbone carries .scales/.biases. The
    model-level sanitize must stack all three, leaving no orphan keys."""
    from omlx.patches.mlx_vlm_mtp import qwen35_moe_vlm_model
    from mlx_vlm.models.qwen3_5_moe import qwen3_5_moe

    monkeypatch.setattr(qwen35_moe_vlm_model, "_APPLIED", False)
    if hasattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched"):
        monkeypatch.delattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched")
    assert qwen35_moe_vlm_model.apply() is True

    pfx_in = "model.language_model.layers.0.mlp"
    weights = {}
    for e in range(2):
        for proj in ("gate_proj", "up_proj", "down_proj"):
            weights[f"{pfx_in}.experts.{e}.{proj}.weight"] = mx.zeros((8, 4))
            weights[f"{pfx_in}.experts.{e}.{proj}.scales"] = mx.zeros((8, 1))
            weights[f"{pfx_in}.experts.{e}.{proj}.biases"] = mx.zeros((8, 1))

    result = qwen3_5_moe.Model.sanitize(_per_expert_vlm_self(), weights)

    pfx = "language_model.model.layers.0.mlp"
    for proj in ("gate_proj", "up_proj", "down_proj"):
        for suffix in ("weight", "scales", "biases"):
            key = f"{pfx}.switch_mlp.{proj}.{suffix}"
            assert key in result, key
            assert result[key].shape[0] == 2
    assert not any(f"{pfx}.experts." in k for k in result)


def test_moe_vlm_sanitize_stacks_per_expert_mtp_quantized(monkeypatch):
    """A per-expert *quantized* MTP head also carries .scales/.biases.
    The model-level VLM sanitize path must keep parity with the runtime
    sanitize path and stack all three suffixes."""
    from omlx.patches.mlx_vlm_mtp import qwen35_moe_vlm_model
    from mlx_vlm.models.qwen3_5_moe import qwen3_5_moe

    monkeypatch.setattr(qwen35_moe_vlm_model, "_APPLIED", False)
    if hasattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched"):
        monkeypatch.delattr(qwen3_5_moe.Model, "_omlx_mtp_vlm_patched")
    assert qwen35_moe_vlm_model.apply() is True

    pfx_in = "mtp.layers.0.mlp"
    weights = {}
    for e in range(2):
        for proj in ("gate_proj", "up_proj", "down_proj"):
            weights[f"{pfx_in}.experts.{e}.{proj}.weight"] = mx.zeros((8, 4))
            weights[f"{pfx_in}.experts.{e}.{proj}.scales"] = mx.zeros((8, 1))
            weights[f"{pfx_in}.experts.{e}.{proj}.biases"] = mx.zeros((8, 1))

    result = qwen3_5_moe.Model.sanitize(_per_expert_vlm_self(), weights)

    pfx = "language_model.mtp.layers.0.mlp"
    for proj in ("gate_proj", "up_proj", "down_proj"):
        for suffix in ("weight", "scales", "biases"):
            key = f"{pfx}.switch_mlp.{proj}.{suffix}"
            assert key in result, key
            assert result[key].shape[0] == 2
    assert not any(f"{pfx}.experts." in k for k in result)


def test_moe_vlm_runtime_sanitize_stacks_per_expert_backbone():
    """The runtime sanitize wrapper must also stack per-expert backbone
    layers (parity with the model-level patch and the LLM patch)."""
    from omlx.patches.mlx_vlm_mtp import qwen35_moe_vlm_runtime

    class FakeModel:
        pass

    fake_outer = SimpleNamespace(Model=FakeModel)
    qwen35_moe_vlm_runtime._patch_vlm_outer_model_sanitize(fake_outer)

    pfx_in = "model.language_model.layers.0.mlp"
    weights = {}
    for e in range(2):
        weights[f"{pfx_in}.experts.{e}.gate_proj.weight"] = mx.zeros((8, 4))
        weights[f"{pfx_in}.experts.{e}.up_proj.weight"] = mx.zeros((8, 4))
        weights[f"{pfx_in}.experts.{e}.down_proj.weight"] = mx.zeros((4, 8))

    result = FakeModel.sanitize(_per_expert_vlm_self(), weights)

    pfx = "language_model.model.layers.0.mlp"
    assert result[f"{pfx}.switch_mlp.gate_proj.weight"].shape == (2, 8, 4)
    assert result[f"{pfx}.switch_mlp.up_proj.weight"].shape == (2, 8, 4)
    assert result[f"{pfx}.switch_mlp.down_proj.weight"].shape == (2, 4, 8)
    assert not any(f"{pfx}.experts." in k for k in result)


# ---------------------------------------------------------------------------
# _call_backbone return format tests
# ---------------------------------------------------------------------------


class TestCallBackbone:
    """Verify _call_backbone handles both tuple and LanguageModelOutput."""

    def test_tuple_2_return(self):
        """mlx-lm dense path returns (logits, hidden) 2-tuple."""
        from omlx.patches.mlx_lm_mtp.batch_generator import _call_backbone

        import mlx.core as mx

        logits = mx.zeros((1, 1, 100))
        hidden = mx.zeros((1, 1, 64))

        model = MagicMock(return_value=(logits, hidden))
        result = _call_backbone(model, mx.zeros((1, 4)), cache=[])
        assert result[0] is logits
        assert result[1] is hidden
        assert result[2] is None  # gdn_states

    def test_tuple_3_return(self):
        """mlx-vlm MoE path returns (logits, hidden, gdn_states) 3-tuple."""
        from omlx.patches.mlx_lm_mtp.batch_generator import _call_backbone

        import mlx.core as mx

        logits = mx.zeros((1, 1, 100))
        hidden = mx.zeros((1, 1, 64))
        gdn = [{"state": "mock"}]

        model = MagicMock(return_value=(logits, hidden, gdn))
        result = _call_backbone(model, mx.zeros((1, 4)), cache=[])
        assert result[0] is logits
        assert result[1] is hidden
        assert result[2] is gdn

    def test_language_model_output_return(self):
        """LanguageModelOutput is correctly unpacked."""
        from omlx.patches.mlx_lm_mtp.batch_generator import _call_backbone

        import mlx.core as mx
        from mlx_vlm.models.base import LanguageModelOutput

        logits = mx.zeros((1, 1, 100))
        hidden = mx.zeros((1, 1, 64))
        gdn = [{"state": "mock"}]

        out = LanguageModelOutput(
            logits=logits,
            hidden_states=[hidden],
            gdn_states=gdn,
        )
        model = MagicMock(return_value=out)
        result = _call_backbone(model, mx.zeros((1, 4)), cache=[])
        assert result[0] is logits
        assert result[1] is hidden
        assert result[2] is gdn