"""Helpers for normalizing worker-advertised and job-requested model identifiers."""

from __future__ import annotations

DEFAULT_MODEL = "large-v3"
DEFAULT_QUANTIZATION = "4bit"
DEFAULT_SUPPORTED_TRANSCRIPTION_MODELS = ["turbo", "large-v3-4bit", "large-v3"]
SUPPORTED_BASE_MODELS = ("large-v3", "turbo")
SUPPORTED_QUANTIZATIONS = ("none", "4bit", "8bit")
TURBO_MODEL_ALIASES = frozenset({"large-v3-turbo", "turbo"})


def supported_model_identifiers(*, default_quantization: str = DEFAULT_QUANTIZATION) -> tuple[str, ...]:
    supported: list[str] = []
    for model in SUPPORTED_BASE_MODELS:
        for quantization in _model_quantizations(model, default_quantization=default_quantization):
            model_id = canonical_model_identifier(
                model,
                quantization,
                default_quantization=default_quantization,
            )
            if model_id not in supported:
                supported.append(model_id)
    return tuple(supported)


def is_supported_model_identifier(
    model: str | None,
    quantization: str | None = None,
    *,
    default_quantization: str = DEFAULT_QUANTIZATION,
) -> bool:
    return canonical_model_identifier(
        model,
        quantization,
        default_quantization=default_quantization,
    ) in supported_model_identifiers(default_quantization=default_quantization)


def _model_quantizations(model: str, *, default_quantization: str) -> tuple[str, ...]:
    if model == DEFAULT_MODEL:
        return (default_quantization, "none", "8bit")
    return ("none",)


def normalize_quantization(value: str | None, *, default: str = "none") -> str:
    if value is None:
        return default
    normalized = str(value).strip().lower()
    return normalized or default


def split_model_identifier(
    model: str | None,
    quantization: str | None = None,
    *,
    default_quantization: str = "none",
) -> tuple[str, str]:
    raw_model = str(model or DEFAULT_MODEL).strip().removeprefix("whisper-")
    inferred_quantization: str | None = None

    if raw_model in TURBO_MODEL_ALIASES:
        base_model = "turbo"
    elif raw_model.endswith("-4bit"):
        base_model = raw_model[:-5]
        inferred_quantization = "4bit"
    elif raw_model.endswith("-8bit"):
        base_model = raw_model[:-5]
        inferred_quantization = "8bit"
    else:
        base_model = raw_model or DEFAULT_MODEL

    if quantization is None:
        if inferred_quantization is not None:
            resolved_quantization = inferred_quantization
        elif base_model == DEFAULT_MODEL:
            resolved_quantization = default_quantization
        else:
            resolved_quantization = "none"
    else:
        resolved_quantization = normalize_quantization(quantization, default="none")
    return base_model, resolved_quantization


def canonical_model_identifier(
    model: str | None,
    quantization: str | None = None,
    *,
    default_quantization: str = "none",
) -> str:
    base_model, resolved_quantization = split_model_identifier(
        model,
        quantization,
        default_quantization=default_quantization,
    )
    if resolved_quantization == "none":
        return base_model
    return f"{base_model}-{resolved_quantization}"
