from typing import Any, Optional

from pydantic import BaseModel, SerializerFunctionWrapHandler, model_serializer
from sqlmodel import Field, Relationship, SQLModel


class AiiProvider(SQLModel, table=True):
    """
    模型提供商。
    """

    id: Optional[int] = Field(default=None, primary_key=True)
    name: str
    base_url: str
    api_key: str

    aii_models: list["AiiModel"] = Relationship(back_populates="aii_provider")

    @model_serializer(mode="wrap")
    def serialize_provider(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]:
        data: dict = handler(self)
        data.pop("api_key", None)
        return data


class AiiProviderPublicWithoutKey(BaseModel):
    id: Optional[int] = None
    name: str
    base_url: str


class AiiProviderPublic(BaseModel):
    id: Optional[int] = None
    name: str
    base_url: str
    api_key: str


class AiiModel(SQLModel, table=True):
    """
    模型。
    """

    id: Optional[int] = Field(default=None, primary_key=True)
    model_name: str
    max_context_length: int
    reasonable: Optional[bool] = Field(default=None, nullable=True, description="模型是否具备思考能力")

    aii_provider_id: int = Field(default=None, foreign_key="aiiprovider.id")
    aii_provider: AiiProvider = Relationship(back_populates="aii_models")

    chatrooms: list["Chatroom"] = Relationship(back_populates="default_model")

    @model_serializer(mode="wrap")
    def serialize_model(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]:
        data: dict = handler(self)
        data["reasonable"] = bool(data.get("reasonable"))
        return data


class AiiModelPublic(BaseModel):
    id: Optional[int] = None
    model_name: str
    max_context_length: int
    reasonable: bool

    aii_provider_id: int


from .model_story import Chatroom  # noqa: E402
