"""Exercise actual attention, convolution and recurrent layers on Metal."""

import mlx.core as mx
import numpy as np
from mlx_vlm.models.qwen3_5.config import TextConfig
from mlx_vlm.models.qwen3_5.language import LanguageModel

from solomon_mlx.engine import Engine, fork_cache


def test_chunked_hybrid_decoder_and_cache_isolation():
    mx.random.seed(8)
    config = TextConfig(
        model_type="qwen3_5_text",
        hidden_size=128,
        intermediate_size=192,
        linear_num_value_heads=2,
        linear_num_key_heads=2,
        linear_key_head_dim=32,
        linear_value_head_dim=32,
        linear_conv_kernel_dim=4,
        num_hidden_layers=4,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=32,
        rms_norm_eps=1e-6,
        vocab_size=256,
        max_position_embeddings=512,
        rope_parameters={
            "type": "default",
            "mrope_section": [2, 1, 1],
            "rope_theta": 100000,
            "partial_rotary_factor": 0.25,
        },
    )
    engine = Engine.__new__(Engine)
    engine.lm = LanguageModel(config)
    engine.lm.eval()

    class ForbiddenVocabularyHead:
        def __call__(self, *args, **kwargs):
            raise AssertionError("Vocabulary projection must never run")

    engine.lm.lm_head = ForbiddenVocabularyHead()
    engine.chunk_size = 16
    engine.context = {"start": None}
    tokens = list(range(1, 74))
    p = 51
    prefix = engine.lm.make_cache()
    engine.forward(tokens[:p], engine.positions(0, p), prefix)
    before = [tuple(None if x is None else np.asarray(x).copy() for x in c.state) for c in prefix]
    a, _ = engine.forward(tokens[p:], engine.positions(p, len(tokens) - p), fork_cache(prefix))
    b, _ = engine.forward(tokens, engine.positions(0, len(tokens)), engine.lm.make_cache())
    np.testing.assert_allclose(np.asarray(a), np.asarray(b), atol=5e-4, rtol=5e-4)
    again, _ = engine.forward(tokens[p:], engine.positions(p, len(tokens) - p), fork_cache(prefix))
    np.testing.assert_array_equal(np.asarray(a), np.asarray(again))
    for original, saved in zip(prefix, before):
        for x, y in zip(original.state, saved):
            if y is not None:
                np.testing.assert_array_equal(np.asarray(x), y)
    for c in prefix:
        if hasattr(c, "cache"):
            assert c.cache[1].dtype == mx.float32