from pydantic import BaseModel, Field
from typing import Any, Dict, List, Optional


class ModelTags(BaseModel):
    modality: List[str] = []
    highlight: List[str] = []
    use_case: List[str] = []
    model_family: List[str] = []

    model_config = {"extra": "allow"}


class ModelRestrictions(BaseModel):
    restriction_type: Optional[str] = None
    restriction_name: Optional[str] = None

    model_config = {"extra": "allow"}


class ServerlessModel(BaseModel):
    model_name: str
    model_alias: Optional[str] = None
    model_type: Optional[str] = None
    context_length: Optional[int] = None
    max_completion_tokens: Optional[int] = None
    description: Optional[str] = None
    input_price: Optional[float] = None
    output_price: Optional[float] = None
    cost_per_1k_input: Optional[float] = None
    cost_per_1k_output: Optional[float] = None
    hosting_in: Optional[str] = None
    parameter_size: Optional[str] = None
    huggingface_url: Optional[str] = None
    readme: Optional[str] = None
    modality: Optional[str] = None
    tier2_plus_require: Optional[bool] = Field(None, alias="tier2+_require")
    restrictions: Optional[Any] = None
    tags: Optional[ModelTags] = None

    model_config = {"extra": "ignore", "populate_by_name": True}


class ServerlessModelsResponse(BaseModel):
    models: List[ServerlessModel]
    count: int
    tag_counts: Optional[Dict[str, Dict[str, int]]] = None
