from pydantic import BaseModel, Field, field_validator
from enum import StrEnum
import datetime as dt


class Modality(StrEnum):
    TEXT = "text"
    IMAGE = "image"
    AUDIO = "audio"
    VIDEO = "video"
    PDF = "pdf"


class ModelModalities(BaseModel):
    input: list[Modality]
    output: list[Modality]


class ModelCost(BaseModel):
    input: float
    output: float


class ModelLimit(BaseModel):
    context: int
    output: int


class Model(BaseModel):
    id: str
    name: str
    family: str | None = None
    attachment: bool
    reasoning: bool
    tool_call: bool
    structured_output: bool | None = None
    temperature: bool | None = None
    knowledge: str | None = None
    release_date: dt.date
    last_updated: dt.date
    modalities: ModelModalities
    open_weights: bool
    cost: ModelCost | None = None
    limit: ModelLimit
    provider_id: str = ""

    @field_validator("release_date", "last_updated", mode="before")
    @classmethod
    def parse_partial_date(cls, v: str) -> str:
        """
        Some models only provide a partial date in the YYYY-MM format, so we
        parse those by defaulting to first day of the month.
        """
        if len(v) == 7:
            return f"{v}-01"
        return v


class Provider(BaseModel):
    id: str
    env: list[str] = Field(default_factory=list)
    npm: str | None = None
    api: str | None = None
    name: str
    doc: str | None = None
    models: dict[str, Model] = Field(default_factory=dict)
