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

# Copyright 2026 The vLLM team.
# Copyright 2026 NVIDIA CORPORATION and the HuggingFace Inc. team. All rights
# reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from importlib.metadata import version
from unittest.mock import MagicMock

import numpy as np
import pytest
import torch
from packaging.version import Version
from transformers import PretrainedConfig

from tests.models.registry import HF_EXAMPLE_MODELS


class MockMusicFlamingoConfig(PretrainedConfig):
    model_type = "musicflamingo"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.audio_config = PretrainedConfig()
        self.text_config = PretrainedConfig()


class MockMusicFlamingoProcessor:
    def __init__(self):
        self.audio_token = "<sound>"
        self.audio_token_id = 12345
        self.audio_bos_token = "<|sound_bos|>"
        self.audio_bos_token_id = 12346
        self.audio_eos_token = "<|sound_eos|>"
        self.audio_eos_token_id = 12347
        self.max_audio_len = 1200
        self.feature_extractor = MockFeatureExtractor()


class MockFeatureExtractor:
    def __init__(self):
        self.sampling_rate = 16000
        self.chunk_length = 30


@pytest.fixture
def mock_ctx():
    config = MockMusicFlamingoConfig()

    ctx = MagicMock()
    ctx.get_hf_config.return_value = config
    ctx.get_hf_processor.return_value = MockMusicFlamingoProcessor()
    ctx.model_config.hf_config = config
    return ctx


@pytest.fixture(autouse=True)
def check_transformers_version():
    model_info = HF_EXAMPLE_MODELS.get_hf_info("MusicFlamingoForConditionalGeneration")
    model_info.check_transformers_version(on_fail="skip")


def test_musicflamingo_chunk_counting_uses_rote_timestamps(mock_ctx, monkeypatch):
    from vllm.model_executor.models.musicflamingo import (
        MusicFlamingoDummyInputsBuilder,
        MusicFlamingoMultiModalProcessor,
        MusicFlamingoProcessingInfo,
    )

    info = MusicFlamingoProcessingInfo(mock_ctx)
    processor = MusicFlamingoMultiModalProcessor(
        info, MusicFlamingoDummyInputsBuilder(info)
    )

    sr = 16000
    audio_1 = np.zeros(30 * sr)
    audio_2 = np.zeros(45 * sr)

    mm_data = {"audio": [audio_1, audio_2]}
    prompt = "<|user|>Listen.<|end|>"

    from vllm.multimodal.processing import BaseMultiModalProcessor

    def mock_base_call(self, prompt, mm_data, mm_kwargs, tok_kwargs):
        del self, prompt, mm_data, mm_kwargs, tok_kwargs
        return {
            "input_ids": [1, 2, 3],
            "input_features": torch.randn(3, 80, 3000),
            "rote_timestamps": torch.randn(3, 750),
        }

    monkeypatch.setattr(BaseMultiModalProcessor, "_call_hf_processor", mock_base_call)

    processed = processor._call_hf_processor(prompt, mm_data, {}, {})

    chunk_counts = processed["chunk_counts"]

    assert chunk_counts.tolist() == [1, 2]
    assert "rote_timestamps" in processed


def test_musicflamingo_dummy_text_uses_plain_audio_tokens(mock_ctx):
    from vllm.model_executor.models.musicflamingo import (
        MusicFlamingoDummyInputsBuilder,
        MusicFlamingoProcessingInfo,
    )

    info = MusicFlamingoProcessingInfo(mock_ctx)
    builder = MusicFlamingoDummyInputsBuilder(info)

    assert builder.get_dummy_text({"audio": 2}) == "<sound><sound>"


@pytest.mark.skipif(
    Version(version("transformers")) >= Version("5.5"),
    reason="transformers v5.5 added native MusicFlamingoForConditionalGeneration "
    "with a different get_audio_features signature (requires input_ids)",
)
def test_musicflamingo_audio_feature_pipeline_matches_hf_small_config():
    from transformers.models.musicflamingo import (
        modeling_musicflamingo as hf_musicflamingo_modeling,
    )
    from transformers.models.musicflamingo.configuration_musicflamingo import (
        MusicFlamingoConfig,
    )

    from vllm.model_executor.models.audioflamingo3 import (
        _build_audio_encoder_attention_mask,
        _flatten_valid_audio_embeddings,
    )
    from vllm.model_executor.models.musicflamingo import (
        MusicFlamingoEncoder,
        MusicFlamingoMultiModalProjector,
        MusicFlamingoRotaryEmbedding,
        apply_rotary_time_emb,
    )

    text_config = {
        "model_type": "qwen2",
        "intermediate_size": 64,
        "initializer_range": 0.02,
        "hidden_size": 32,
        "max_position_embeddings": 1024,
        "num_hidden_layers": 2,
        "num_attention_heads": 4,
        "num_key_value_heads": 2,
        "vocab_size": 128,
        "pad_token_id": 1,
        "use_mrope": False,
    }
    audio_config = {
        "hidden_size": 16,
        "num_attention_heads": 4,
        "intermediate_size": 32,
        "num_hidden_layers": 2,
        "num_mel_bins": 80,
        "max_source_positions": 1500,
        "dropout": 0.0,
        "attention_dropout": 0.0,
        "activation_dropout": 0.0,
        "encoder_layerdrop": 0.0,
    }

    torch.manual_seed(0)
    config = MusicFlamingoConfig(
        text_config=text_config,
        audio_config=audio_config,
        audio_token_id=0,
        head_dim=8,
        rope_parameters={"rope_type": "default", "rope_theta": 2048},
    )
    hf_model = hf_musicflamingo_modeling.MusicFlamingoForConditionalGeneration(
        config
    ).eval()

    vllm_encoder = MusicFlamingoEncoder(config.audio_config).eval()
    vllm_encoder.load_state_dict(hf_model.audio_tower.state_dict())

    vllm_projector = MusicFlamingoMultiModalProjector(config).eval()
    vllm_projector.load_state_dict(hf_model.multi_modal_projector.state_dict())

    vllm_rope = MusicFlamingoRotaryEmbedding(config).eval()
    vllm_rope.load_state_dict(hf_model.pos_emb.state_dict(), strict=False)

    input_features = torch.randn(3, 80, 3000)
    feature_attention_mask = torch.zeros(3, 3000, dtype=torch.bool)
    feature_attention_mask[0, :3000] = True
    feature_attention_mask[1, :2500] = True
    feature_attention_mask[2, :1500] = True
    rote_timestamps = (
        torch.arange(750, dtype=torch.float32).unsqueeze(0).repeat(3, 1) * 0.04
    )

    hf_output = hf_model.get_audio_features(
        input_features,
        feature_attention_mask,
        rote_timestamps=rote_timestamps,
        return_dict=True,
    ).pooler_output
    vllm_attention_mask = _build_audio_encoder_attention_mask(
        feature_attention_mask,
        dtype=vllm_encoder.conv1.weight.dtype,
        device=vllm_encoder.conv1.weight.device,
    )
    vllm_hidden_states = vllm_encoder(
        input_features,
        attention_mask=vllm_attention_mask,
    )
    cos, sin = vllm_rope(rote_timestamps, seq_len=vllm_hidden_states.shape[-2])
    vllm_hidden_states = apply_rotary_time_emb(vllm_hidden_states, cos, sin)
    vllm_output, _ = _flatten_valid_audio_embeddings(
        vllm_projector(vllm_hidden_states),
        feature_attention_mask,
    )

    torch.testing.assert_close(vllm_output, hf_output)