"""Pydantic models for custom AI model configurations."""

from __future__ import annotations

from pydantic import BaseModel, ConfigDict, Field

from ..aimodelscommon import SchemaUrl
from ..aimodelsdatamodels import BaseModelConfigSchema, ModelPack
from ..aimodelsproviders.ai_models_providers import ModelProviderExtraAuthVars
from ..core.enums import ModelPublisher


def _empty_base_model_config_schema_list() -> list[BaseModelConfigSchema]:
    return []


def _empty_custom_model_list() -> list[CustomModel]:
    return []


def _empty_custom_provider_list() -> list[CustomProvider]:
    return []


def _empty_model_pack_list() -> list[ModelPack]:
    return []


def _empty_client_model_pack_schema_list() -> list[ClientModelPackSchema]:
    return []


__all__ = [
    "ClientModelPackSchema",
    "ClientModelsInput",
    "CustomAIModel",
    "CustomModel",
    "CustomProvider",
    "ModelsInput",
]


class CustomModel(BaseModel):
    """Represents a custom AI model configuration."""

    id: str | None = None
    model_id: str = Field(..., alias="modelId")
    publisher: ModelPublisher
    description: str
    providers: list[BaseModelConfigSchema] = Field(
        default_factory=_empty_base_model_config_schema_list
    )
    created_at: str | None = Field(None, alias="createdAt")
    updated_at: str | None = Field(None, alias="updatedAt")

    model_config = ConfigDict(
        populate_by_name=True,
        str_strip_whitespace=True,
        validate_assignment=True,
    )


class CustomProvider(BaseModel):
    """Represents a custom AI model provider."""

    id: str | None = None
    name: str
    base_url: str = Field(..., alias="baseUrl")
    has_aws_auth: bool = Field(False, alias="hasAWSAuth")
    skip_auth: bool = Field(False, alias="skipAuth")
    api_key_env_var: str | None = Field(None, alias="apiKeyEnvVar")
    extra_auth_vars: list[ModelProviderExtraAuthVars] | None = Field(
        None, alias="extraAuthVars"
    )
    created_at: str | None = Field(None, alias="createdAt")
    updated_at: str | None = Field(None, alias="updatedAt")

    model_config = ConfigDict(
        populate_by_name=True,
        str_strip_whitespace=True,
        validate_assignment=True,
    )


class ClientModelPackSchema(BaseModel):
    """Schema for a client-side model pack."""

    name: str
    description: str


class ModelsInput(BaseModel):
    """Input for creating or updating models, providers, and model packs."""

    models: list[CustomModel] = Field(default_factory=_empty_custom_model_list)
    providers: list[CustomProvider] = Field(default_factory=_empty_custom_provider_list)
    model_packs: list[ModelPack] = Field(
        default_factory=_empty_model_pack_list, alias="modelPacks"
    )

    model_config = ConfigDict(populate_by_name=True)


class ClientModelsInput(BaseModel):
    """Input for client-side model configurations."""

    schema_url: SchemaUrl | None = Field(None, alias="$schema")
    models: list[CustomModel] = Field(default_factory=_empty_custom_model_list)
    providers: list[CustomProvider] = Field(default_factory=_empty_custom_provider_list)
    model_packs: list[ClientModelPackSchema] = Field(
        default_factory=_empty_client_model_pack_schema_list, alias="modelPacks"
    )

    model_config = ConfigDict(populate_by_name=True)


class CustomAIModel(BaseModel):
    """Data contract for CustomAIModel."""

    provider_id: str = Field(alias="providerId")
    model_id: str = Field(alias="modelId")
    is_chat: bool = Field(default=False, alias="isChat")
    is_image_to_text: bool = Field(default=False, alias="isImageToText")
    is_multi_modal: bool = Field(default=False, alias="isMultiModal")
    is_fine_tunable: bool = Field(default=False, alias="isFineTunable")
    is_function_calling: bool = Field(default=False, alias="isFunctionCalling")

    model_config = ConfigDict(populate_by_name=True)
