from lalamo.model_import.decoder_configs import HFQwen2Config, HFQwen3Config
from lalamo.quantization import QuantizationMode

from .common import ConfigMap, FileSpec, ModelSpec, UseCase, WeightsType

__all__ = ["QWEN_MODELS"]


QWEN25 = [
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5",
        name="Qwen2.5-0.5B-Instruct",
        size="0.5B",
        quantization=None,
        repo="Qwen/Qwen2.5-0.5B-Instruct",
        config_type=HFQwen2Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5",
        name="Qwen2.5-1.5B-Instruct",
        size="1.5B",
        quantization=None,
        repo="Qwen/Qwen2.5-1.5B-Instruct",
        config_type=HFQwen2Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5",
        name="Qwen2.5-3B-Instruct",
        size="3B",
        quantization=None,
        repo="Qwen/Qwen2.5-3B-Instruct",
        config_type=HFQwen2Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5",
        name="Qwen2.5-7B-Instruct",
        size="7B",
        quantization=None,
        repo="Qwen/Qwen2.5-7B-Instruct",
        config_type=HFQwen2Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5",
        name="Qwen2.5-14B-Instruct",
        size="14B",
        quantization=None,
        repo="Qwen/Qwen2.5-14B-Instruct",
        config_type=HFQwen2Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5",
        name="Qwen2.5-32B-Instruct",
        size="32B",
        quantization=None,
        repo="Qwen/Qwen2.5-32B-Instruct",
        config_type=HFQwen2Config,
    ),
]


QWEN25_CODER = [
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5-Coder",
        name="Qwen2.5-Coder-0.5B-Instruct",
        size="0.5B",
        quantization=None,
        repo="Qwen/Qwen2.5-Coder-0.5B-Instruct",
        config_type=HFQwen2Config,
        use_cases=(UseCase.CODE,),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5-Coder",
        name="Qwen2.5-Coder-1.5B-Instruct",
        size="1.5B",
        quantization=None,
        repo="Qwen/Qwen2.5-Coder-1.5B-Instruct",
        config_type=HFQwen2Config,
        use_cases=(UseCase.CODE,),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5-Coder",
        name="Qwen2.5-Coder-3B-Instruct",
        size="3B",
        quantization=None,
        repo="Qwen/Qwen2.5-Coder-3B-Instruct",
        config_type=HFQwen2Config,
        use_cases=(UseCase.CODE,),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5-Coder",
        name="Qwen2.5-Coder-7B-Instruct",
        size="7B",
        quantization=None,
        repo="Qwen/Qwen2.5-Coder-7B-Instruct",
        config_type=HFQwen2Config,
        use_cases=(UseCase.CODE,),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5-Coder",
        name="Qwen2.5-Coder-14B-Instruct",
        size="14B",
        quantization=None,
        repo="Qwen/Qwen2.5-Coder-14B-Instruct",
        config_type=HFQwen2Config,
        use_cases=(UseCase.CODE,),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5-Coder",
        name="Qwen2.5-Coder-32B-Instruct",
        size="32B",
        quantization=None,
        repo="Qwen/Qwen2.5-Coder-32B-Instruct",
        config_type=HFQwen2Config,
        use_cases=(UseCase.CODE,),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen2.5-Coder",
        name="Qwen2.5-Coder-32B-Instruct",
        size="32B",
        quantization=None,
        repo="Qwen/Qwen2.5-Coder-32B-Instruct",
        config_type=HFQwen2Config,
        use_cases=(UseCase.CODE,),
    ),
]


QWEN3 = [
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-0.6B",
        size="0.6B",
        quantization=None,
        repo="Qwen/Qwen3-0.6B",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-0.6B-MLX-4bit",
        size="0.6B",
        quantization=QuantizationMode.UINT4,
        repo="Qwen/Qwen3-0.6B-MLX-4bit",
        config_type=HFQwen3Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-0.6B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-0.6B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-0.6B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-1.7B",
        size="1.7B",
        quantization=None,
        repo="Qwen/Qwen3-1.7B",
        config_type=HFQwen3Config,
        weights_type=WeightsType.SAFETENSORS,
        use_cases=tuple(),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-4B",
        size="4B",
        quantization=None,
        repo="Qwen/Qwen3-4B",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-4B-AWQ",
        size="4B",
        quantization=QuantizationMode.UINT4,
        repo="Qwen/Qwen3-4B-AWQ",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-4B-MLX-4bit",
        size="4B",
        quantization=QuantizationMode.UINT4,
        repo="Qwen/Qwen3-4B-MLX-4bit",
        config_type=HFQwen3Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-4B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-4B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-4B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-8B",
        size="8B",
        quantization=None,
        repo="Qwen/Qwen3-8B",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-8B-AWQ",
        size="8B",
        quantization=QuantizationMode.UINT4,
        repo="Qwen/Qwen3-8B-AWQ",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-8B-MLX-4bit",
        size="8B",
        quantization=QuantizationMode.UINT4,
        repo="Qwen/Qwen3-8B-MLX-4bit",
        config_type=HFQwen3Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-8B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-8B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-8B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-14B",
        size="14B",
        quantization=None,
        repo="Qwen/Qwen3-14B",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-14B-AWQ",
        size="14B",
        quantization=None,
        repo="Qwen/Qwen3-14B-AWQ",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-32B",
        size="32B",
        quantization=None,
        repo="Qwen/Qwen3-32B",
        config_type=HFQwen3Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3",
        name="Qwen3-32B-AWQ",
        size="32B",
        quantization=QuantizationMode.UINT4,
        repo="Qwen/Qwen3-32B-AWQ",
        config_type=HFQwen3Config,
    ),
]


QWEN_MODELS = QWEN25 + QWEN25_CODER + QWEN3
