from lalamo.model_import.model_configs import (
    HFBonsaiConfig,
    HFQwen2Config,
    HFQwen3Config,
    HFQwen3NextConfig,
    HFQwen35Config,
)
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,
    ),
]

QWEN3_NEXT = [
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3-Next",
        name="Qwen3-Next-80B-A3B-Instruct",
        size="80B",
        quantization=None,
        repo="Qwen/Qwen3-Next-80B-A3B-Instruct",
        config_type=HFQwen3NextConfig,
    ),
    # TODO(knyazer): uncomment and test when sharding is in
    # ModelSpec(
    #     vendor="Alibaba",
    #     family="Qwen3-Next",
    #     name="Qwen3-Next-80B-A3B-Instruct-4bit",
    #     size="80B",
    #     quantization=QuantizationMode.UINT4,
    #     repo="mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit",
    #     config_type=HFQwen3NextConfig,
    #     configs=ConfigMap(
    #         tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3-Next-80B-A3B-Instruct"),
    #         tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3-Next-80B-A3B-Instruct"),
    #         generation_config=FileSpec("generation_config.json", "Qwen/Qwen3-Next-80B-A3B-Instruct"),
    #     ),
    # ),
]

QWEN35 = [
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-0.8B",
        size="0.8B",
        quantization=None,
        repo="Qwen/Qwen3.5-0.8B",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-0.8B-MLX-4bit",
        size="0.8B",
        quantization=QuantizationMode.UINT4,
        repo="mlx-community/Qwen3.5-0.8B-MLX-4bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-0.8B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-0.8B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-0.8B-MLX-8bit",
        size="0.8B",
        quantization=QuantizationMode.UINT8,
        repo="mlx-community/Qwen3.5-0.8B-MLX-8bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-0.8B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-0.8B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-2B",
        size="2B",
        quantization=None,
        repo="Qwen/Qwen3.5-2B",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-2B-MLX-4bit",
        size="2B",
        quantization=QuantizationMode.UINT4,
        repo="mlx-community/Qwen3.5-2B-MLX-4bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-2B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-2B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-2B-MLX-8bit",
        size="2B",
        quantization=QuantizationMode.UINT8,
        repo="mlx-community/Qwen3.5-2B-MLX-8bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-2B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-2B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-4B",
        size="4B",
        quantization=None,
        repo="Qwen/Qwen3.5-4B",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-4B-MLX-4bit",
        size="4B",
        quantization=QuantizationMode.UINT4,
        repo="mlx-community/Qwen3.5-4B-MLX-4bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-4B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-4B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-4B-MLX-8bit",
        size="4B",
        quantization=QuantizationMode.UINT8,
        repo="mlx-community/Qwen3.5-4B-MLX-8bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-4B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-4B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-9B",
        size="9B",
        quantization=None,
        repo="Qwen/Qwen3.5-9B",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-9B-MLX-4bit",
        size="9B",
        quantization=QuantizationMode.UINT4,
        repo="mlx-community/Qwen3.5-9B-MLX-4bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-9B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-9B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-9B-MLX-8bit",
        size="9B",
        quantization=QuantizationMode.UINT8,
        repo="mlx-community/Qwen3.5-9B-MLX-8bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-9B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-9B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-27B",
        size="27B",
        quantization=None,
        repo="Qwen/Qwen3.5-27B",
        config_type=HFQwen35Config,
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-27B-MLX-4bit",
        size="27B",
        quantization=QuantizationMode.UINT4,
        repo="mlx-community/Qwen3.5-27B-4bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-27B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-27B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
    ModelSpec(
        vendor="Alibaba",
        family="Qwen3.5",
        name="Qwen3.5-27B-MLX-8bit",
        size="27B",
        quantization=QuantizationMode.UINT8,
        repo="mlx-community/Qwen3.5-27B-8bit",
        config_type=HFQwen35Config,
        configs=ConfigMap(
            tokenizer=FileSpec("tokenizer.json", "Qwen/Qwen3.5-27B"),
            tokenizer_config=FileSpec("tokenizer_config.json", "Qwen/Qwen3.5-27B"),
            generation_config=FileSpec("generation_config.json", "Qwen/Qwen3.5-27B"),
        ),
    ),
]

QWEN3_PARD = [
    ModelSpec(
        vendor="AMD",
        family="PARD",
        name="PARD-Qwen3-0.6B",
        size="0.6B",
        quantization=None,
        repo="amd/PARD-Qwen3-0.6B",
        config_type=HFQwen3Config,
    ),
]

BONSAI = [
    ModelSpec(
        vendor="PrismML",
        family="Bonsai",
        name="Bonsai-4B-mlx-1bit",
        size="4B",
        quantization=QuantizationMode.UINT1,
        repo="prism-ml/Bonsai-4B-mlx-1bit",
        config_type=HFBonsaiConfig,
        configs=ConfigMap(
            chat_template=FileSpec("chat_template.jinja"),
        ),
    ),
]

QWEN_MODELS = QWEN25 + QWEN25_CODER + QWEN3 + QWEN3_NEXT + QWEN35 + QWEN3_PARD + BONSAI
