# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from types import SimpleNamespace

import numpy as np
import pytest
import torch
from transformers import Qwen3Config

from vllm.model_executor.models.interfaces import SupportsLoRA, supports_lora
from vllm.model_executor.models.moss_audio import (
    MOSS_AUDIO_BOS_TOKEN,
    MOSS_AUDIO_BOS_TOKEN_ID,
    MOSS_AUDIO_EOS_TOKEN,
    MOSS_AUDIO_EOS_TOKEN_ID,
    MOSS_AUDIO_PLACEHOLDER,
    MOSS_AUDIO_TOKEN,
    MOSS_AUDIO_TOKEN_ID,
    MossAudioConfig,
    MossAudioDummyInputsBuilder,
    MossAudioEncoder,
    MossAudioEncoderConfig,
    MossAudioModel,
    MossAudioMultiModalProcessor,
    MossAudioProcessingInfo,
    MossAudioProcessor,
    MossQwen3ForCausalLM,
    MossQwen3Model,
)
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
from vllm.multimodal.inputs import batched_tensors_equal
from vllm.sequence import IntermediateTensors


class _Tokenizer:
    def encode(self, text, add_special_tokens=False):
        del add_special_tokens
        return [ord(char) for char in text]

    def decode(self, token_ids, **kwargs):
        del kwargs
        return "".join(chr(token_id) for token_id in token_ids)

    def batch_decode(self, batch_token_ids, **kwargs):
        return [self.decode(token_ids, **kwargs) for token_ids in batch_token_ids]


class _MMConfig:
    enable_mm_embeds = False
    mm_processor_cache_gb = 1

    def merge_mm_processor_kwargs(self, kwargs):
        return dict(kwargs)

    def get_limit_per_prompt(self, modality):
        del modality
        return 3


class _ModelConfig:
    def __init__(self):
        self.model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
        self.revision = None
        self.max_model_len = 4096
        self.encoder_config = {}
        self.dtype = torch.float32
        self.hf_config = MossAudioConfig(language_config=Qwen3Config())
        self.multimodal_config = _MMConfig()

    def get_multimodal_config(self):
        return self.multimodal_config

    def get_inputs_embeds_size(self):
        return None


class _ProcessingContext:
    def __init__(self):
        self.model_config = _ModelConfig()
        self.tokenizer = _Tokenizer()

    def get_tokenizer(self):
        return self.tokenizer

    def get_hf_config(self):
        return self.model_config.hf_config

    def get_mm_config(self):
        return self.model_config.get_multimodal_config()

    def get_merged_mm_kwargs(self, kwargs):
        return self.get_mm_config().merge_mm_processor_kwargs(kwargs)

    def call_hf_processor(self, hf_processor, data, kwargs):
        merged_kwargs = self.get_merged_mm_kwargs(kwargs)
        merged_kwargs.setdefault("return_tensors", "pt")
        return hf_processor(**data, **merged_kwargs)


class _TestMossAudioProcessingInfo(MossAudioProcessingInfo):
    def _get_processor_config_defaults(self):
        return {}


def _vllm_config(tensor_parallel_size=1, pipeline_parallel_size=1, hf_config=None):
    if hf_config is None:
        hf_config = MossAudioConfig(language_config=Qwen3Config())
    return SimpleNamespace(
        model_config=SimpleNamespace(
            hf_config=hf_config,
            multimodal_config=None,
        ),
        quant_config=None,
        parallel_config=SimpleNamespace(
            tensor_parallel_size=tensor_parallel_size,
            pipeline_parallel_size=pipeline_parallel_size,
        ),
    )


class _FakeAudioEncoder:
    dtype = torch.float32

    def __init__(self, deepstack_layers=0):
        self.deepstack_layers = deepstack_layers
        self.output_deepstack_hidden_states = None
        self.input_shape = None
        self.feature_lens = None

    def __call__(self, audio_data, *, feature_lens, output_deepstack_hidden_states):
        self.input_shape = tuple(audio_data.shape)
        self.feature_lens = feature_lens.detach().cpu().clone()
        self.output_deepstack_hidden_states = output_deepstack_hidden_states
        lengths = MossAudioEncoder._compute_downsampled_length(feature_lens)
        hidden_states = torch.ones(1, int(lengths.sum().item()), 8)
        if not output_deepstack_hidden_states:
            return hidden_states, None
        return hidden_states, [
            hidden_states * scale for scale in range(2, 2 + self.deepstack_layers)
        ]


def _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=1, tp_rank=0):
    import vllm.model_executor.layers.linear as linear_layers
    import vllm.model_executor.models.moss_audio as moss_audio_module
    import vllm.model_executor.parameter as parameter_module

    for module in (moss_audio_module, linear_layers, parameter_module):
        monkeypatch.setattr(
            module, "get_tensor_model_parallel_world_size", lambda: tp_size
        )
    monkeypatch.setattr(
        linear_layers, "get_tensor_model_parallel_rank", lambda: tp_rank
    )
    monkeypatch.setattr(
        parameter_module, "get_tensor_model_parallel_rank", lambda: tp_rank
    )
    monkeypatch.setattr(
        linear_layers, "tensor_model_parallel_all_reduce", lambda tensor: tensor
    )


def _build_moss_audio_processor(cache=None):
    ctx = _ProcessingContext()
    info = _TestMossAudioProcessingInfo(ctx)
    return (
        MossAudioMultiModalProcessor(
            info,
            MossAudioDummyInputsBuilder(info),
            cache=cache,
        ),
        ctx,
    )


def _assert_mm_inputs_equal(left, right):
    assert left["prompt_token_ids"] == right["prompt_token_ids"]
    assert left["mm_hashes"] == right["mm_hashes"]

    left_placeholder = left["mm_placeholders"]["audio"][0]
    right_placeholder = right["mm_placeholders"]["audio"][0]
    assert left_placeholder.offset == right_placeholder.offset
    assert left_placeholder.length == right_placeholder.length
    assert left_placeholder.is_embed.tolist() == right_placeholder.is_embed.tolist()

    assert batched_tensors_equal(
        left["mm_kwargs"].get_data(),
        right["mm_kwargs"].get_data(),
    )


@pytest.mark.parametrize(
    ("prompt", "prefix"),
    [
        (
            f"before {MOSS_AUDIO_PLACEHOLDER} after",
            [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
        ),
        (
            f"before {MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}"
            f"{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN} after",
            [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
        ),
        ("Describe this audio.", [MOSS_AUDIO_BOS_TOKEN_ID]),
    ],
)
def test_moss_audio_processor_expands_audio_placeholders(prompt, prefix):
    raw_mel_len = 17
    processed = MossAudioProcessor(_Tokenizer())(
        text=prompt, audio=[torch.zeros(160 * raw_mel_len)]
    )
    input_ids = processed["input_ids"][0].tolist()

    assert input_ids[: len(prefix)] == prefix
    assert input_ids.count(MOSS_AUDIO_BOS_TOKEN_ID) == 1
    assert input_ids.count(MOSS_AUDIO_EOS_TOKEN_ID) == 1
    assert input_ids.count(MOSS_AUDIO_TOKEN_ID) == (
        MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
    )
    assert processed["audio_data"].shape == (1, 128, raw_mel_len)
    assert processed["audio_data_seqlens"].tolist() == [raw_mel_len]


def test_moss_audio_processor_preserves_placeholder_without_audio():
    processed = MossAudioProcessor(_Tokenizer())(
        text=f"before {MOSS_AUDIO_PLACEHOLDER} after"
    )

    assert processed["input_ids"][0].tolist() == [
        *[ord(char) for char in "before "],
        MOSS_AUDIO_BOS_TOKEN_ID,
        MOSS_AUDIO_TOKEN_ID,
        MOSS_AUDIO_EOS_TOKEN_ID,
        *[ord(char) for char in " after"],
    ]
    assert "audio_data" not in processed
    assert "audio_data_seqlens" not in processed


def test_moss_audio_multimodal_processor_handles_token_and_cache_paths():
    raw_mel_len = 17
    audio = np.zeros(160 * raw_mel_len, dtype=np.float32)
    prompt = f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."

    baseline_processor, ctx = _build_moss_audio_processor()
    mm_items = baseline_processor.info.parse_mm_data({"audio": [audio]})
    token_prompt = ctx.get_tokenizer().encode(prompt, add_special_tokens=False)

    baseline_text = baseline_processor(
        prompt,
        mm_items=mm_items,
        hf_processor_mm_kwargs={},
    )
    baseline_token = baseline_processor(
        token_prompt,
        mm_items=mm_items,
        hf_processor_mm_kwargs={},
    )

    cache = MultiModalProcessorOnlyCache(ctx.model_config)
    cached_processor, _ = _build_moss_audio_processor(cache=cache)
    cached_text_miss = cached_processor(
        prompt,
        mm_items=mm_items,
        hf_processor_mm_kwargs={},
    )
    cached_text_hit = cached_processor(
        prompt,
        mm_items=mm_items,
        hf_processor_mm_kwargs={},
    )
    cached_token_hit = cached_processor(
        token_prompt,
        mm_items=mm_items,
        hf_processor_mm_kwargs={},
    )

    expected_audio_tokens = MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
    prompt_token_ids = baseline_text["prompt_token_ids"]
    assert prompt_token_ids.count(MOSS_AUDIO_TOKEN_ID) == expected_audio_tokens
    assert baseline_text["mm_placeholders"]["audio"][0].length == (
        expected_audio_tokens + 2
    )

    _assert_mm_inputs_equal(baseline_text, baseline_token)
    _assert_mm_inputs_equal(baseline_text, cached_text_miss)
    _assert_mm_inputs_equal(baseline_text, cached_text_hit)
    _assert_mm_inputs_equal(baseline_text, cached_token_hit)


def test_moss_audio_supports_language_model_lora_only():
    assert supports_lora(MossAudioModel)

    model = object.__new__(MossAudioModel)
    assert isinstance(model, SupportsLoRA)

    mapping = model.get_mm_mapping()
    assert mapping.language_model == ["language_model."]
    assert mapping.tower_model == []
    assert mapping.connector == []


def test_moss_audio_error_paths():
    model = object.__new__(MossAudioModel)
    with pytest.raises(ValueError, match="DeepStack audio token count mismatch"):
        model._cache_deepstack_input_embeds(
            inputs_embeds=torch.zeros(4, 8),
            deepstack_embeddings=((torch.ones(1, 8),),),
            is_multimodal=torch.tensor([False, True, True, False]),
        )

    with pytest.raises(ValueError, match="too short"):
        MossAudioProcessor(_Tokenizer())(
            text=MOSS_AUDIO_PLACEHOLDER, audio=[torch.empty(0)]
        )
    with pytest.raises(ValueError, match="too short"):
        model._parse_and_validate_audio_input(
            audio_data=torch.zeros(1, 128, 1),
            audio_data_seqlens=torch.tensor([0], dtype=torch.long),
        )


def test_moss_audio_validates_tp_config():
    vllm_config = _vllm_config(tensor_parallel_size=2)
    vllm_config.model_config.hf_config.adapter_hidden_size = 7

    with pytest.raises(ValueError, match="adapter_hidden_size"):
        MossAudioModel(vllm_config=vllm_config)

    vllm_config = _vllm_config(tensor_parallel_size=2)
    vllm_config.model_config.hf_config.audio_config.d_model = 6
    vllm_config.model_config.hf_config.audio_config.encoder_attention_heads = 3
    with pytest.raises(ValueError, match="encoder_attention_heads"):
        MossAudioModel(vllm_config=vllm_config)


def test_moss_audio_rejects_audio_data_list_seqlen_count_mismatch():
    model = object.__new__(MossAudioModel)

    with pytest.raises(ValueError, match="audio_data batch size"):
        model._parse_and_validate_audio_input(
            audio_data=[torch.zeros(128, 8), torch.zeros(128, 11)],
            audio_data_seqlens=torch.tensor([8], dtype=torch.long),
        )


@pytest.mark.parametrize("deepstack_scales", [(), (7, 11)])
def test_moss_audio_embed_multimodal_packs_by_audio(deepstack_scales):
    model = object.__new__(MossAudioModel)
    model.audio_encoder = _FakeAudioEncoder(len(deepstack_scales))
    model.audio_adapter = lambda hidden_states: hidden_states * 5
    model.deepstack_audio_merger_list = [
        lambda hidden_states, scale=scale: hidden_states * scale
        for scale in deepstack_scales
    ]
    model.deepstack_input_embeds = None

    embeddings = model.embed_multimodal(
        audio_data=torch.zeros(2, 128, 9),
        audio_data_seqlens=torch.tensor([8, 9], dtype=torch.long),
    )

    assert model.audio_encoder.output_deepstack_hidden_states is bool(deepstack_scales)
    assert [embeds.shape for embeds in embeddings] == [
        torch.Size([1, 8 * (1 + len(deepstack_scales))]),
        torch.Size([2, 8 * (1 + len(deepstack_scales))]),
    ]
    if not deepstack_scales:
        assert model.deepstack_input_embeds is None
        return

    main_embeddings, deepstack_embeddings = model._split_multimodal_embeddings(
        embeddings, hidden_size=8
    )
    assert [embeds.shape for embeds in main_embeddings] == [
        torch.Size([1, 8]),
        torch.Size([2, 8]),
    ]
    assert [[e.shape for e in layer] for layer in deepstack_embeddings] == [
        [torch.Size([1, 8]), torch.Size([2, 8])] for _ in deepstack_scales
    ]
    assert torch.equal(main_embeddings[0], torch.full((1, 8), 5.0))
    for idx, scale in enumerate(deepstack_scales):
        assert torch.equal(
            deepstack_embeddings[idx][0],
            torch.full((1, 8), float((idx + 2) * scale)),
        )


def test_moss_audio_embed_input_ids_caches_packed_deepstack():
    class _FakeLanguageModel:
        def embed_input_ids(self, input_ids):
            return torch.zeros(input_ids.shape[0], 8)

    model = object.__new__(MossAudioModel)
    model.language_model = _FakeLanguageModel()
    model.deepstack_audio_merger_list = [object(), object()]
    model.deepstack_input_embeds = None
    multimodal_embeddings = (
        torch.cat([torch.full((1, 8), x) for x in (5.0, 14.0, 33.0)], dim=-1),
        torch.cat([torch.full((2, 8), x) for x in (7.0, 22.0, 44.0)], dim=-1),
    )
    is_multimodal = torch.tensor([False, True, True, True, False])

    inputs_embeds = model.embed_input_ids(
        input_ids=torch.arange(5),
        multimodal_embeddings=multimodal_embeddings,
        is_multimodal=is_multimodal,
    )

    assert torch.equal(inputs_embeds[1], torch.full((8,), 5.0))
    assert torch.equal(inputs_embeds[2], torch.full((8,), 7.0))
    assert torch.equal(inputs_embeds[3], torch.full((8,), 7.0))
    assert model.deepstack_input_embeds is not None
    tensors = model.deepstack_input_embeds.tensors
    assert set(tensors) == {"deepstack_input_embeds_0", "deepstack_input_embeds_1"}
    for tensor in tensors.values():
        assert tensor[is_multimodal].abs().sum() > 0
        assert torch.equal(tensor[~is_multimodal], torch.zeros(2, 8))


def _patch_pp_group(monkeypatch, *, first=True, last=True):
    import vllm.model_executor.models.moss_audio as moss_audio_module

    monkeypatch.setattr(
        moss_audio_module,
        "get_pp_group",
        lambda: SimpleNamespace(is_first_rank=first, is_last_rank=last),
    )


def test_moss_audio_pp_forward_routes_deepstack(monkeypatch):
    for first in (True, False):
        calls: list[dict[str, object]] = []

        def fake_lm(*args, _calls=calls, **kwargs):
            del args
            _calls.append(kwargs)
            return torch.ones(1, 1)

        _patch_pp_group(monkeypatch, first=first)
        model = object.__new__(MossAudioModel)
        torch.nn.Module.__init__(model)
        model.language_model = fake_lm
        cached = IntermediateTensors({"deepstack_input_embeds_0": torch.ones(3, 8)})
        inter = IntermediateTensors(
            {
                "hidden_states": torch.ones(3, 8),
                "residual": torch.zeros(3, 8),
                "deepstack_input_embeds_0": torch.full((3, 8), 5.0),
            }
        )
        inputs_embeds = torch.full((3, 8), 9.0)
        model.deepstack_input_embeds = cached

        model.forward(
            input_ids=None,
            positions=torch.arange(3),
            intermediate_tensors=None if first else inter,
            inputs_embeds=inputs_embeds if first else None,
        )

        kwargs = calls[0]
        assert kwargs["inputs_embeds"] is (inputs_embeds if first else None)
        assert kwargs["deepstack_input_embeds"] is (cached if first else inter)
        assert model.deepstack_input_embeds is None

    calls = []

    def fake_lm_non_first_rank(*args, **kwargs):
        del args
        calls.append(kwargs)
        return torch.ones(1, 1)

    _patch_pp_group(monkeypatch, first=False)
    model = object.__new__(MossAudioModel)
    torch.nn.Module.__init__(model)
    model.language_model = fake_lm_non_first_rank
    model.deepstack_input_embeds = IntermediateTensors({})
    inter = IntermediateTensors(
        {
            "hidden_states": torch.ones(3, 8),
            "residual": torch.zeros(3, 8),
        }
    )

    model.forward(
        input_ids=None,
        positions=torch.arange(3),
        intermediate_tensors=inter,
        inputs_embeds=torch.ones(3, 8),
    )
    assert calls[0]["inputs_embeds"] is None
    assert calls[0]["deepstack_input_embeds"] is inter


def test_moss_qwen3_deepstack_keys_for_pp(monkeypatch):
    class AddOne(torch.nn.Module):
        def forward(self, positions, hidden_states, residual):
            del positions, residual
            return hidden_states + 1, torch.zeros_like(hidden_states)

    def make_model(num_layers, deepstack_layers=None):
        model = object.__new__(MossQwen3Model)
        torch.nn.Module.__init__(model)
        model.start_layer, model.end_layer = 0, num_layers
        model.layers = torch.nn.ModuleList([AddOne() for _ in range(num_layers)])
        model.norm = lambda hidden_states, residual: (hidden_states, residual)
        model._maybe_add_hidden_state = lambda aux, *args: aux
        model.deepstack_inject_layer_indices = (
            range(0) if deepstack_layers is None else deepstack_layers
        )
        return model

    _patch_pp_group(monkeypatch, first=True, last=True)
    output = make_model(3).forward(
        input_ids=None,
        positions=torch.arange(2),
        inputs_embeds=torch.zeros(2, 4),
        deepstack_input_embeds=IntermediateTensors(
            {
                "deepstack_input_embeds_2": torch.full((2, 4), 5.0),
            }
        ),
    )
    assert torch.equal(output, torch.full((2, 4), 8.0))

    _patch_pp_group(monkeypatch, first=True, last=False)
    deepstack = IntermediateTensors(
        {
            "deepstack_input_embeds_0": torch.full((2, 4), 7.0),
            "deepstack_input_embeds_3": torch.full((2, 4), 11.0),
        }
    )
    output = make_model(2, range(4)).forward(
        input_ids=None,
        positions=torch.arange(2),
        inputs_embeds=torch.zeros(2, 4),
        deepstack_input_embeds=deepstack,
    )
    assert isinstance(output, IntermediateTensors)
    assert set(output.tensors) == {
        "hidden_states",
        "residual",
        "deepstack_input_embeds_2",
        "deepstack_input_embeds_3",
    }
    assert torch.equal(output["hidden_states"], torch.full((2, 4), 9.0))
    assert torch.equal(output["deepstack_input_embeds_2"], torch.zeros(2, 4))
    assert output["deepstack_input_embeds_3"] is deepstack["deepstack_input_embeds_3"]

    inner_model = make_model(0, range(2))
    inner_model.make_empty_intermediate_tensors = lambda batch, dtype, device: (
        IntermediateTensors(
            {
                "hidden_states": torch.zeros(batch, 4, dtype=dtype, device=device),
                "residual": torch.zeros(batch, 4, dtype=dtype, device=device),
            }
        )
    )
    language_model = object.__new__(MossQwen3ForCausalLM)
    torch.nn.Module.__init__(language_model)
    language_model.model = inner_model
    language_model.config = SimpleNamespace(hidden_size=4)
    language_model.deepstack_inject_layer_indices = range(2)

    tensors = MossQwen3ForCausalLM.make_empty_intermediate_tensors(
        language_model,
        batch_size=3,
        dtype=torch.float16,
        device=torch.device("cpu"),
    )

    assert set(tensors.tensors) == {
        "hidden_states",
        "residual",
        "deepstack_input_embeds_0",
        "deepstack_input_embeds_1",
    }
    assert tensors["deepstack_input_embeds_0"].shape == (3, 4)
    assert tensors["deepstack_input_embeds_0"].dtype == torch.float16

    _patch_pp_group(monkeypatch, first=True, last=False)
    forward_tensors = inner_model.forward(
        input_ids=None,
        positions=torch.arange(3),
        inputs_embeds=torch.ones(3, 4, dtype=torch.float16),
        deepstack_input_embeds=None,
    )
    assert isinstance(forward_tensors, IntermediateTensors)
    assert set(forward_tensors.tensors) == set(tensors.tensors)


def test_moss_audio_encoder_loads_realistic_attention_weight_names(monkeypatch):
    from vllm.config import VllmConfig, set_current_vllm_config
    from vllm.config.device import DeviceConfig

    _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=2)
    config = MossAudioEncoderConfig(
        d_model=8,
        output_dim=8,
        num_mel_bins=8,
        encoder_layers=1,
        encoder_attention_heads=2,
        encoder_ffn_dim=16,
        downsample_hidden_size=2,
        deepstack_encoder_layer_indexes=[],
    )
    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
        encoder = MossAudioEncoder(config)

    attention = encoder.layers[0].self_attn
    assert all(hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj"))
    assert hasattr(attention, "out_proj")
    assert not hasattr(attention, "qkv")
    assert attention.k_proj.bias is None

    weight_names = [
        "layers.0.self_attn.q_proj.weight",
        "layers.0.self_attn.q_proj.bias",
        "layers.0.self_attn.k_proj.weight",
        "layers.0.self_attn.v_proj.weight",
        "layers.0.self_attn.v_proj.bias",
        "layers.0.self_attn.out_proj.weight",
        "layers.0.self_attn.out_proj.bias",
        "conv1.weight",
        "conv1.bias",
    ]
    params = dict(encoder.named_parameters(remove_duplicate=False))
    assert "layers.0.self_attn.k_proj.bias" not in params
    weights = {
        name: torch.full_like(params[name], fill_value=float(i + 1))
        for i, name in enumerate(weight_names)
    }

    loaded = AutoWeightsLoader(encoder).load_weights(weights.items())

    assert "load_weights" not in MossAudioEncoder.__dict__
    assert loaded == set(weight_names)
    assert not any(".qkv." in name for name in loaded)
    assert torch.equal(
        params["layers.0.self_attn.q_proj.weight"],
        weights["layers.0.self_attn.q_proj.weight"],
    )
