# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for component routing for quantization."""

from unittest.mock import MagicMock

import pytest
import torch
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig,
)
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
    CompressedTensorsConfig,
)
from vllm.model_executor.models.utils import WeightsMapper

from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
    PRE_QUANTIZED_METHODS,
)
from vllm_omni.quantization.component_config import (
    ComponentQuantizationConfig,
    resolve_component_quant_config,
)
from vllm_omni.quantization.inc_config import OmniINCConfig

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


# ---------------------------------------------------------------------------
# Helpers: lightweight mock quant configs
# ---------------------------------------------------------------------------


class _MockQuantConfig(QuantizationConfig):
    """Minimal mock that only implements get_name()."""

    def __init__(self, name: str, **attrs):
        self._name = name
        for k, v in attrs.items():
            setattr(self, k, v)

    def get_name(self) -> str:
        return self._name

    def get_quant_method(self, layer, prefix):
        return MagicMock()

    @classmethod
    def get_supported_act_dtypes(cls):
        return [torch.bfloat16, torch.float16]

    def get_min_capability(self):
        return 0

    @classmethod
    def from_config(cls, config):
        raise NotImplementedError

    def get_config_filenames(self):
        return []


def _make_inc_config(block_names="thinker.model.layers,talker.model.layers", extra_config=None):
    """Create a real OmniINCConfig with block_name_to_quantize."""
    return OmniINCConfig(
        weight_bits=4,
        group_size=128,
        sym=True,
        block_name_to_quantize=block_names,
        extra_config=extra_config or {},
    )


THINKER_MAPPER = WeightsMapper(
    orig_to_new_prefix={
        "thinker.lm_head.": "language_model.lm_head.",
        "thinker.model.": "language_model.model.",
        "thinker.": "",
    }
)

TALKER_MAPPER = WeightsMapper(
    orig_to_new_prefix={
        "talker.codec_head.": "language_model.lm_head.",
        "talker.model.": "language_model.model.",
        "talker.thinker_to_talker_proj.": "thinker_to_talker_proj.",
        "talker.": "",
    }
)


# ===================================================================
# 1. OmniINCConfig.apply_vllm_mapper
# ===================================================================


class TestApplyVllmMapper:
    def test_inc_csv_string_normalized_to_list(self):
        """CSV string block_name_to_quantize is split into a list."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        assert isinstance(cfg.block_name_to_quantize, list)

    def test_thinker_blocks_remapped(self):
        """thinker.model.layers -> language_model.model.layers after apply_vllm_mapper."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize)

    def test_cross_stage_blocks_kept_unchanged(self):
        """Blocks not matching any mapper prefix are kept unchanged (harmless)."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        # talker.model.layers doesn't match any thinker mapper prefix → stays as-is
        assert "talker.model.layers" in cfg.block_name_to_quantize

    def test_talker_remap(self):
        """talker.model.layers -> language_model.model.layers with talker mapper."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(TALKER_MAPPER)
        assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize)
        # thinker.model.layers doesn't match talker mapper → stays as-is
        assert "thinker.model.layers" in cfg.block_name_to_quantize

    def test_extra_config_keys_remapped(self):
        """Regex keys in extra_config get their escaped-dot prefixes remapped."""
        extra = {
            r".*thinker\.model\.layers\.0\.mlp\.gate.*": {"bits": 16, "data_type": "float"},
        }
        cfg = _make_inc_config("thinker.model.layers", extra_config=extra)
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        # The key should now reference the vLLM runtime path
        assert any("language_model" in k for k in cfg.extra_config)
        # Original thinker\.model prefix should be replaced
        assert not any(r"thinker\.model" in k for k in cfg.extra_config)

    def test_single_block_name(self):
        """Only one block name (not CSV) still works."""
        cfg = _make_inc_config("thinker.model.layers")
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize)

    def test_already_list_block_names(self):
        """block_name_to_quantize already a list (not CSV string) works."""
        cfg = _make_inc_config(["thinker.model.layers", "talker.model.layers"])
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        assert isinstance(cfg.block_name_to_quantize, list)
        assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize)

    def test_mutates_in_place(self):
        """apply_vllm_mapper mutates the config in place (same as upstream INCConfig)."""
        cfg = _make_inc_config("thinker.model.layers")
        original_id = id(cfg)
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        assert id(cfg) == original_id

    # -- Stage prefix tests (runtime prefix = container + internal name) --

    def test_thinker_block_has_stage_prefix(self):
        """Mapped block name must start with 'thinker.' so runtime startswith() works."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        assert "thinker.language_model.model.layers" in cfg.block_name_to_quantize

    def test_talker_block_has_stage_prefix(self):
        """Mapped block name must start with 'talker.' so runtime startswith() works."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(TALKER_MAPPER)
        assert "talker.language_model.model.layers" in cfg.block_name_to_quantize

    def test_thinker_block_matches_runtime_prefix(self):
        """Simulates get_layer_config's startswith() check for FusedMoE layers."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        runtime_prefix = "thinker.language_model.model.layers.0.mlp.experts"
        assert any(runtime_prefix.startswith(b) for b in cfg.block_name_to_quantize)

    def test_talker_block_matches_runtime_prefix(self):
        """Simulates get_layer_config's startswith() check for talker FusedMoE."""
        cfg = _make_inc_config("thinker.model.layers,talker.model.layers")
        cfg.apply_vllm_mapper(TALKER_MAPPER)
        runtime_prefix = "talker.language_model.model.layers.0.mlp.experts"
        assert any(runtime_prefix.startswith(b) for b in cfg.block_name_to_quantize)

    def test_extra_config_plain_key_has_stage_prefix(self):
        """Plain extra_config keys are remapped with stage prefix."""
        extra = {
            "talker.model.layers.0.mlp.shared_expert_gate": {"bits": 16},
        }
        cfg = _make_inc_config("talker.model.layers", extra_config=extra)
        cfg.apply_vllm_mapper(TALKER_MAPPER)
        assert "talker.language_model.model.layers.0.mlp.shared_expert_gate" in cfg.extra_config

    def test_extra_config_regex_key_still_works(self):
        """Regex extra_config keys use re.search so no stage prefix needed."""
        import re

        extra = {
            r".*thinker\.model\.layers\.0\.mlp\.gate.*": {"bits": 16},
        }
        cfg = _make_inc_config("thinker.model.layers", extra_config=extra)
        cfg.apply_vllm_mapper(THINKER_MAPPER)
        runtime_name = "thinker.language_model.model.layers.0.mlp.gate"
        matched = any(re.search(k, runtime_name) for k in cfg.extra_config)
        assert matched


# ===================================================================
# 2. OmniINCConfig upgrade helpers
# ===================================================================


class TestOmniINCConfigUpgrade:
    def test_maybe_upgrade_none(self):
        assert OmniINCConfig.maybe_upgrade(None) is None

    def test_maybe_upgrade_non_inc(self):
        """Non-INC configs are passed through unchanged."""
        cfg = _MockQuantConfig("fp8")
        assert OmniINCConfig.maybe_upgrade(cfg) is cfg

    def test_maybe_upgrade_already_omni(self):
        """Already OmniINCConfig is returned as-is."""
        cfg = _make_inc_config()
        assert OmniINCConfig.maybe_upgrade(cfg) is cfg

    def test_maybe_upgrade_vanilla_inc(self):
        """Vanilla INCConfig is promoted to OmniINCConfig."""
        from vllm.model_executor.layers.quantization.inc import INCConfig

        vanilla = INCConfig(weight_bits=4, group_size=128, sym=True)
        upgraded = OmniINCConfig.maybe_upgrade(vanilla)
        assert isinstance(upgraded, OmniINCConfig)
        assert upgraded.weight_bits == 4
        assert upgraded.group_size == 128


# ===================================================================
# 2. Three-branch thinker routing (simulated)
# ===================================================================


def _simulate_thinker_routing(quant_config):
    """Simulate the three-branch routing in thinker __init__.

    Returns (visual_quant_config, language_quant_config, wrapped_vllm_quant).
    """
    if isinstance(quant_config, ComponentQuantizationConfig):
        visual_quant_config = quant_config.resolve("visual")
        language_quant_config = quant_config.resolve("language_model")
        return visual_quant_config, language_quant_config, quant_config
    elif quant_config is not None:
        if quant_config.get_name() in PRE_QUANTIZED_METHODS:
            return quant_config, quant_config, quant_config
        else:
            language_quant_config = quant_config
            wrapped = ComponentQuantizationConfig(
                component_configs={"language_model": quant_config},
                default_config=None,
            )
            return None, language_quant_config, wrapped
    else:
        return None, None, None


class TestThinkerRouting:
    def test_none(self):
        vis, lang, wrapped = _simulate_thinker_routing(None)
        assert vis is None
        assert lang is None
        assert wrapped is None

    @pytest.mark.parametrize("method", ["modelopt", "modelopt_fp4", "modelopt_mxfp8"])
    def test_pre_quantized_all_components(self, method):
        """Pre-quantized methods pass config to all components."""
        cfg = _MockQuantConfig(method)
        vis, lang, wrapped = _simulate_thinker_routing(cfg)
        assert vis is cfg
        assert lang is cfg
        assert wrapped is cfg

    def test_fp8_dynamic_language_only(self):
        """fp8 dynamic: visual=None, language gets original config."""
        cfg = _MockQuantConfig("fp8")
        vis, lang, wrapped = _simulate_thinker_routing(cfg)
        assert vis is None
        assert lang is cfg
        assert isinstance(wrapped, ComponentQuantizationConfig)
        assert wrapped.resolve("language_model") is cfg
        assert wrapped.resolve("visual") is None

    def test_inc_autoround_language_only(self):
        """INC/AutoRound: not in _PRE_QUANTIZED_METHODS -> wrapped like fp8."""
        cfg = _MockQuantConfig("inc")
        vis, lang, wrapped = _simulate_thinker_routing(cfg)
        assert vis is None
        assert lang is cfg
        assert isinstance(wrapped, ComponentQuantizationConfig)

    def test_component_config_passthrough(self):
        """Explicit ComponentQuantizationConfig is used directly."""
        inner_fp8 = _MockQuantConfig("fp8")
        inner_modelopt = _MockQuantConfig("modelopt")
        cqc = ComponentQuantizationConfig(
            component_configs={
                "visual": inner_modelopt,
                "language_model": inner_fp8,
            }
        )
        vis, lang, wrapped = _simulate_thinker_routing(cqc)
        assert vis is inner_modelopt
        assert lang is inner_fp8
        assert wrapped is cqc


# ===================================================================
# 3. Talker visual routing (init_multi_modal guard)
# ===================================================================


def _simulate_talker_visual_routing(quant_config):
    """Simulate the talker init_multi_modal visual routing."""
    if quant_config is not None and quant_config.get_name() in PRE_QUANTIZED_METHODS:
        return quant_config
    return None


class TestTalkerVisualRouting:
    def test_none(self):
        assert _simulate_talker_visual_routing(None) is None

    @pytest.mark.parametrize("method", ["modelopt", "modelopt_fp4", "modelopt_mxfp8"])
    def test_pre_quantized_passes_through(self, method):
        """Pre-quantized methods pass quant config to visual."""
        cfg = _MockQuantConfig(method)
        assert _simulate_talker_visual_routing(cfg) is cfg

    def test_fp8_blocked(self):
        """fp8 dynamic must NOT be passed to visual."""
        cfg = _MockQuantConfig("fp8")
        assert _simulate_talker_visual_routing(cfg) is None

    def test_inc_blocked(self):
        """INC/AutoRound must NOT be passed to visual (not in _PRE_QUANTIZED_METHODS)."""
        cfg = _MockQuantConfig("inc")
        assert _simulate_talker_visual_routing(cfg) is None


# ===================================================================
# 4. ComponentQuantizationConfig.apply_vllm_mapper
# ===================================================================


class TestComponentApplyVllmMapper:
    def test_maps_compressed_tensors_component_and_default(self):
        thinker = CompressedTensorsConfig(
            target_scheme_map={},
            ignore=["thinker.model.layers.0.mlp.gate"],
            quant_format="pack-quantized",
        )
        default = CompressedTensorsConfig(
            target_scheme_map={},
            ignore=["talker.model.layers.0.mlp.gate"],
            quant_format="pack-quantized",
        )
        config = ComponentQuantizationConfig(
            component_configs={"thinker": thinker, "code2wav": None},
            default_config=default,
        )
        mapper = WeightsMapper(
            orig_to_new_prefix={
                "thinker.model": "thinker.language_model.model",
                "talker.model": "talker.language_model.model",
            }
        ).get_unstacked_mapper()

        config.apply_vllm_mapper(mapper)

        assert thinker.ignore == ["thinker.language_model.model.layers.0.mlp.gate"]
        assert default.ignore == ["talker.language_model.model.layers.0.mlp.gate"]


# ===================================================================
# 5. ComponentQuantizationConfig.resolve
# ===================================================================


class TestResolveComponentQuantConfig:
    def test_global_config_applies_to_each_quantization_aware_component(self):
        config = _MockQuantConfig("fp8")

        assert resolve_component_quant_config(config, "transformer") is config
        assert resolve_component_quant_config(config, "text_encoder") is config

    def test_component_map_is_the_only_scope_override(self):
        transformer = _MockQuantConfig("fp8")
        config = ComponentQuantizationConfig(
            component_configs={"transformer": transformer},
            default_config=None,
        )

        assert resolve_component_quant_config(config, "transformer") is transformer
        assert resolve_component_quant_config(config, "text_encoder") is None


class TestComponentResolve:
    def test_longest_prefix_match(self):
        a = _MockQuantConfig("a")
        b = _MockQuantConfig("b")
        cqc = ComponentQuantizationConfig(component_configs={"language_model": a, "language_model.model": b})
        assert cqc.resolve("language_model.model.layers.0") is b
        assert cqc.resolve("language_model.lm_head") is a

    def test_no_match_returns_default(self):
        a = _MockQuantConfig("a")
        default = _MockQuantConfig("default")
        cqc = ComponentQuantizationConfig(
            component_configs={"language_model": a},
            default_config=default,
        )
        assert cqc.resolve("visual") is default

    def test_no_match_no_default_returns_none(self):
        a = _MockQuantConfig("a")
        cqc = ComponentQuantizationConfig(
            component_configs={"language_model": a},
        )
        assert cqc.resolve("visual") is None

    def test_get_name(self):
        cqc = ComponentQuantizationConfig(component_configs={})
        assert cqc.get_name() == "component"

    def test_get_quant_method_delegates(self):
        """get_quant_method dispatches to the resolved config."""
        inner = _MockQuantConfig("fp8")
        cqc = ComponentQuantizationConfig(
            component_configs={"language_model": inner},
        )
        layer = MagicMock()
        result = cqc.get_quant_method(layer, "language_model.model.layers.0.mlp")
        assert result is not None  # delegates to inner.get_quant_method

    def test_get_quant_method_returns_none_for_unmatched(self):
        """get_quant_method returns None when no config matches."""
        inner = _MockQuantConfig("fp8")
        cqc = ComponentQuantizationConfig(
            component_configs={"language_model": inner},
        )
        layer = MagicMock()
        result = cqc.get_quant_method(layer, "visual.blocks.0.mlp")
        assert result is None

    def test_min_capability(self):
        a = _MockQuantConfig("a")
        a.get_min_capability = lambda: 80
        b = _MockQuantConfig("b")
        b.get_min_capability = lambda: 70
        cqc = ComponentQuantizationConfig(component_configs={"x": a, "y": b})
        assert cqc.get_min_capability() == 70

    def test_min_capability_empty(self):
        cqc = ComponentQuantizationConfig(component_configs={})
        assert cqc.get_min_capability() == 0