# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Tests for the quantization config factory's method-name and disk-dict resolution."""

from __future__ import annotations

import pytest
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import (
    CompressedTensorsConfig,
)

from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
from vllm_omni.quantization.component_config import (
    ComponentQuantizationConfig,
    is_prequantized_config,
    resolve_encoder_quant_config,
)
from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS, _build_single


def _int8_weight_only_disk_config() -> dict:
    """Serialized compressed-tensors int8 weight-only metadata as it appears on disk."""
    weights = {
        "num_bits": 8,
        "type": "int",
        "symmetric": True,
        "strategy": "channel",
        "dynamic": False,
        "observer": "minmax",
    }
    return {
        "quant_method": "compressed-tensors",
        "format": "pack-quantized",
        "config_groups": {"group_0": {"targets": ["Linear"], "weights": weights}},
        "ignore": ["img_in"],
        "quantization_status": "frozen",
    }


def test_hyphenated_registry_method_builds_from_disk_schema() -> None:
    # A bare str spec cannot construct compressed-tensors (a serialized-format
    # method with no parameterless default); the serialized schema must build.
    with pytest.raises(TypeError, match="Cannot instantiate CompressedTensorsConfig"):
        _build_single("compressed-tensors")
    disk = _int8_weight_only_disk_config()
    disk.pop("quant_method")
    config = _build_single("compressed-tensors", **disk)
    assert isinstance(config, CompressedTensorsConfig)


@pytest.mark.parametrize("spelling", ["compressed-tensors", "compressed_tensors"])
def test_disk_dict_builds_via_from_config(spelling: str) -> None:
    disk = _int8_weight_only_disk_config()
    disk["quant_method"] = spelling
    config = build_quant_config(disk)
    assert isinstance(config, CompressedTensorsConfig)


def test_per_component_dict_routes_transformer_only() -> None:
    disk = {"transformer": _int8_weight_only_disk_config()}
    config = build_quant_config(disk)
    assert isinstance(config, ComponentQuantizationConfig)
    assert isinstance(
        resolve_component_quant_config(config, "transformer"), CompressedTensorsConfig
    )
    encoder = resolve_encoder_quant_config(
        resolve_component_quant_config(config, "text_encoder")
    )
    assert encoder is None


def test_unknown_method_still_raises() -> None:
    with pytest.raises(ValueError, match="Unknown quantization method"):
        _build_single("not-a-real-method")


def test_supported_list_advertises_hyphenated_registry_names() -> None:
    assert "compressed-tensors" in SUPPORTED_QUANTIZATION_METHODS


def test_is_prequantized_config_recognizes_serialized_formats() -> None:
    # compressed-tensors is always a serialized format: its linear methods
    # consume packed checkpoint tensors and there is no on-load quantization.
    assert is_prequantized_config(build_quant_config(_int8_weight_only_disk_config()))
    component = build_quant_config({"transformer": _int8_weight_only_disk_config()})
    assert isinstance(component, ComponentQuantizationConfig)
    assert is_prequantized_config(component)
    assert not is_prequantized_config(None)
    # Online (non-serialized) methods stay unrecognized so the loader keeps
    # the accelerator-side load-quantize-offload path for them.
    from vllm_omni.quantization.fp8_config import DiffusionFp8Config

    assert not is_prequantized_config(DiffusionFp8Config())