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


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

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

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


class AiiProviderPublic(BaseModel):
    id: int | None = None
    name: str
    base_url: str
    api_key: str


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

    id: int | None = Field(default=None, primary_key=True)
    model_name: str
    max_context_length: int

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


class AiiModelPublic(BaseModel):
    id: int | None = None
    model_name: str
    max_context_length: int

    aii_provider_id: int


def z_aii_model(am: AiiModel) -> dict:
    return {
        "id": am.id,
        "model_name": am.model_name,
        "max_context_length": am.max_context_length,
        "aii_provider_id": am.aii_provider_id,
    }


def z_aii_provider(ap: AiiProvider) -> dict:
    return {
        "id": ap.id,
        "name": ap.name,
        "base_url": ap.base_url,
    }
