import mlx.core as mx

from ark_asr_mlx.config import ArkASRConfig
from ark_asr_mlx.model import ArkASRModel


def _tiny_config() -> ArkASRConfig:
    return ArkASRConfig.from_dict(
        {
            "hidden_size": 8,
            "intermediate_size": 16,
            "num_attention_heads": 2,
            "num_hidden_layers": 2,
            "num_key_value_heads": 1,
            "rms_norm_eps": 1e-6,
            "rope_theta": 10_000,
            "vocab_size": 20,
            "audio_token_id": 10,
            "eos_token_id": 2,
            "pad_token_id": 0,
            "merge_factor": 2,
            "whisper_config": {
                "d_model": 8,
                "encoder_attention_heads": 2,
                "encoder_ffn_dim": 16,
                "encoder_layers": 1,
                "num_mel_bins": 4,
            },
        }
    )


def test_prompt_injection_and_kv_cache() -> None:
    model = ArkASRModel(_tiny_config())
    input_ids = mx.array([[1, 10, 10, 3, 4]])
    input_features = mx.zeros((1, 4, 8))
    embeddings = model.prepare_prompt_embeddings(
        input_ids,
        input_features,
        audio_start=1,
        audio_count=2,
    )
    assert embeddings.shape == (1, 5, 8)

    cache = model.make_cache()
    prompt_logits = model(input_ids, cache=cache, input_embeddings=embeddings)
    next_logits = model(mx.array([[5]]), cache=cache)
    mx.eval(prompt_logits, next_logits)

    assert prompt_logits.shape == (1, 5, 20)
    assert next_logits.shape == (1, 1, 20)
    assert all(layer_cache.offset == 6 for layer_cache in cache)