from typing import Literal, get_args

from fastvideo.layers.quantization.base_config import QuantizationConfig

QuantizationMethods = Literal[None, "AbsMaxFP8", "FP8", "NVFP4", "nvfp4_qat", "nvfp4_qat_train", "fp8_qat_train"]

QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))

# The customized quantization methods which will be added to this dict.
_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}


def register_quantization_config(quantization: str):
    """Register a customized vllm quantization config.

    When a quantization method is not supported by vllm, you can register a customized
    quantization config to support it.

    Args:
        quantization (str): The quantization method name.

    Examples:
        >>> from fastvideo.layers.quantization import register_quantization_config
        >>> from fastvideo.layers.quantization import get_quantization_config
        >>> from fastvideo.layers.quantization.base_config import QuantizationConfig
        >>>
        >>> @register_quantization_config("my_quant")
        ... class MyQuantConfig(QuantizationConfig):
        ...     pass
        >>>
        >>> get_quantization_config("my_quant")
        <class 'MyQuantConfig'>
    """  # noqa: E501

    def _wrapper(quant_config_cls):
        if quantization in QUANTIZATION_METHODS:
            raise ValueError(f"The quantization method `{quantization}` is already exists.")
        if not issubclass(quant_config_cls, QuantizationConfig):
            raise ValueError("The quantization config must be a subclass of "
                             "`QuantizationConfig`.")
        _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
        QUANTIZATION_METHODS.append(quantization)
        return quant_config_cls

    return _wrapper


def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
    if quantization not in QUANTIZATION_METHODS:
        raise ValueError(f"Invalid quantization method: {quantization}")

    # lazy import to avoid triggering `torch.compile` too early
    from .absmax_fp8 import AbsMaxFP8Config
    from .fp8_config import FP8Config
    from .nvfp4_config import NVFP4Config
    from .nvfp4_qat_config import NVFP4QATConfig
    from .nvfp4_qat_train_config import NVFP4QATTrainConfig
    from .fp8_qat_train_config import FP8QATTrainConfig

    method_to_config: dict[str, type[QuantizationConfig]] = {
        "AbsMaxFP8": AbsMaxFP8Config,
        "FP8": FP8Config,
        "NVFP4": NVFP4Config,
        "nvfp4_qat": NVFP4QATConfig,
        "nvfp4_qat_train": NVFP4QATTrainConfig,
        "fp8_qat_train": FP8QATTrainConfig,
    }
    # Update the `method_to_config` with customized quantization methods.
    method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)

    return method_to_config[quantization]


all = [
    "QuantizationMethods",
    "QuantizationConfig",
    "get_quantization_config",
    "QUANTIZATION_METHODS",
]