from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, Field


class ModelDefinition(BaseModel):
    name: str
    type: Literal["LLM", "TTS", "STS", "VLM"] = "LLM"
    security_rank: int
    interface_cls: str
    args: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)
    default_quant: str | None = None
    quant_files: dict[str, str] = Field(default_factory=dict)

    def build_kwargs(self, **kwargs) -> dict:
        new_quant = kwargs.pop("quant", self.default_quant)

        if new_quant and len(self.quant_files) == 0:
            raise ValueError(
                f"Cannot set a quant for {self.name}. "
                f"Its definition has no model files - is this an API-based model?"
            )

        if new_quant and new_quant not in self.quant_files:
            available = ", ".join(self.quant_files.keys())
            raise ValueError(
                f"The quant '{new_quant}' was not found for model {self.name}. "
                f"Available quants are {available}."
            )

        kws = {**self.args, **kwargs}

        if new_quant:
            kws["model_path"] = Path(self.quant_files[new_quant])

        return kws
