"""Supported model catalog and automatic model selection."""

from __future__ import annotations

from dataclasses import dataclass
import platform


MLX_RUNTIME = "mlx"
LLAMA_CPP_RUNTIME = "llama.cpp"


@dataclass(frozen=True)
class ModelSpec:
    repo_id: str
    filename: str | None
    label: str
    param_count: int
    min_memory_gb: float
    runtime: str = LLAMA_CPP_RUNTIME
    notes: str = ""


APPLE_DEFAULT_MODEL_REPO = "mlx-community/DeepSeek-R1-Distill-Qwen-7B-4bit-mlx"
GENERIC_DEFAULT_MODEL_REPO = "bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF"


SUPPORTED_MODELS: tuple[ModelSpec, ...] = (
    ModelSpec(
        repo_id="mlx-community/DeepSeek-R1-Distill-Qwen-1.5B-4bit",
        filename=None,
        label="DeepSeek R1 Distill Qwen 1.5B 4-bit",
        param_count=1_500_000_000,
        min_memory_gb=4,
        runtime=MLX_RUNTIME,
        notes="Smallest Apple Silicon tier for constrained hosts.",
    ),
    ModelSpec(
        repo_id="bartowski/DeepSeek-R1-Distill-Qwen-1.5B-GGUF",
        filename="DeepSeek-R1-Distill-Qwen-1.5B-Q4_K_M.gguf",
        label="DeepSeek R1 Distill Qwen 1.5B Q4_K_M",
        param_count=1_500_000_000,
        min_memory_gb=4,
        notes="Smallest cross-platform tier with broad llama.cpp compatibility.",
    ),
    ModelSpec(
        repo_id="mlx-community/DeepSeek-R1-Distill-Qwen-7B-4bit-mlx",
        filename=None,
        label="DeepSeek R1 Distill Qwen 7B 4-bit",
        param_count=7_000_000_000,
        min_memory_gb=8,
        runtime=MLX_RUNTIME,
        notes="Balanced Apple Silicon default tier for reasoning and code generation.",
    ),
    ModelSpec(
        repo_id="bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF",
        filename="DeepSeek-R1-Distill-Qwen-7B-Q4_K_M.gguf",
        label="DeepSeek R1 Distill Qwen 7B Q4_K_M",
        param_count=7_000_000_000,
        min_memory_gb=8,
        notes="Balanced cross-platform default tier for reasoning and code generation.",
    ),
    ModelSpec(
        repo_id="mlx-community/DeepSeek-R1-Distill-Qwen-14B-4bit",
        filename=None,
        label="DeepSeek R1 Distill Qwen 14B 4-bit",
        param_count=14_000_000_000,
        min_memory_gb=16,
        runtime=MLX_RUNTIME,
        notes="Higher-quality Apple Silicon tier for larger-memory hosts.",
    ),
    ModelSpec(
        repo_id="bartowski/DeepSeek-R1-Distill-Qwen-14B-GGUF",
        filename="DeepSeek-R1-Distill-Qwen-14B-Q4_K_M.gguf",
        label="DeepSeek R1 Distill Qwen 14B Q4_K_M",
        param_count=14_000_000_000,
        min_memory_gb=16,
        notes="Higher-quality cross-platform tier for larger-memory hosts.",
    ),
    ModelSpec(
        repo_id="mlx-community/DeepSeek-R1-Distill-Qwen-32B-4bit",
        filename=None,
        label="DeepSeek R1 Distill Qwen 32B 4-bit",
        param_count=32_000_000_000,
        min_memory_gb=28,
        runtime=MLX_RUNTIME,
        notes="Largest Apple Silicon tier for high-memory machines.",
    ),
    ModelSpec(
        repo_id="bartowski/DeepSeek-R1-Distill-Qwen-32B-GGUF",
        filename="DeepSeek-R1-Distill-Qwen-32B-Q4_K_M.gguf",
        label="DeepSeek R1 Distill Qwen 32B Q4_K_M",
        param_count=32_000_000_000,
        min_memory_gb=28,
        notes="Largest cross-platform tier for high-memory machines.",
    ),
)


def choose_best_model(available_memory_gb: float | None) -> ModelSpec:
    """Return the largest model that fits, preferring the native runtime on ties."""
    if available_memory_gb is None:
        return get_default_model_spec()

    fitting_models = [
        model for model in SUPPORTED_MODELS if available_memory_gb >= model.min_memory_gb
    ]
    if not fitting_models:
        return get_default_model_spec()
    return max(fitting_models, key=_model_sort_key)


def get_default_model_repo() -> str:
    """Return the default repo id for the current host."""
    if platform.system() == "Darwin" and platform.machine() == "arm64":
        return APPLE_DEFAULT_MODEL_REPO
    return GENERIC_DEFAULT_MODEL_REPO


def get_default_model_spec() -> ModelSpec:
    """Return the project default provider model."""
    for model in SUPPORTED_MODELS:
        if model.repo_id == get_default_model_repo():
            return model
    return SUPPORTED_MODELS[0]


def get_model_spec(model_name: str) -> ModelSpec | None:
    """Return the model spec for a known repo id."""
    for model in SUPPORTED_MODELS:
        if model.repo_id == model_name:
            return model
    return None


def get_model_param_count(model_name: str) -> int:
    """Return parameter count for a known model, else default to 7B."""
    model = get_model_spec(model_name)
    return model.param_count if model else 7_000_000_000


def _model_sort_key(model: ModelSpec) -> tuple[int, int]:
    """Prefer the largest fitting tier, then the platform-native runtime."""
    prefers_mlx = platform.system() == "Darwin" and platform.machine() == "arm64"
    native_runtime_score = int(
        model.runtime == (MLX_RUNTIME if prefers_mlx else LLAMA_CPP_RUNTIME)
    )
    return (model.param_count, native_runtime_score)
