"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

from __future__ import annotations
from enum import Enum
import httpx
from leonardo_ai_sdk.types import (
    BaseModel,
    Nullable,
    OptionalNullable,
    UNSET,
    UNSET_SENTINEL,
)
import pydantic
from pydantic import model_serializer
from typing import Optional
from typing_extensions import Annotated, NotRequired, TypedDict


class SdVersions(str, Enum):
    r"""The base version of stable diffusion to use if not using a custom model."""

    SDXL_0_9 = "SDXL_0_9"
    SDXL_1_0 = "SDXL_1_0"
    LEONARDO_DIFFUSION_XL = "LEONARDO_DIFFUSION_XL"
    LEONARDO_LIGHTNING_XL = "LEONARDO_LIGHTNING_XL"
    VISION_XL = "VISION_XL"
    KINO_XL = "KINO_XL"
    ALBEDO_XL = "ALBEDO_XL"


class CreateElementRequestBodyTypedDict(TypedDict):
    r"""Query parameters to be provided in the request body as a JSON object."""

    dataset_id: NotRequired[str]
    r"""The ID of the dataset to train the element on."""
    description: NotRequired[Nullable[str]]
    r"""The description of the element."""
    instance_prompt: NotRequired[str]
    r"""The instance prompt to use during training.Try “a” by a noun. E.g. a castle"""
    learning_rate: NotRequired[float]
    r"""The speed of element learns."""
    lora_focus: NotRequired[str]
    r"""The category determines how the element will be trained. Options are 'General' | 'Character' | 'Style' | 'Object'."""
    name: NotRequired[str]
    r"""The name of the element."""
    num_train_epochs: NotRequired[int]
    r"""The number of times the entire training dataset is passed through the element."""
    resolution: NotRequired[Nullable[int]]
    r"""The resolution for training. Must be 1024."""
    sd_version: NotRequired[SdVersions]
    r"""The base version of stable diffusion to use if not using a custom model."""
    train_text_encoder: NotRequired[bool]
    r"""Whether or not encode the train text."""


class CreateElementRequestBody(BaseModel):
    r"""Query parameters to be provided in the request body as a JSON object."""

    dataset_id: Annotated[Optional[str], pydantic.Field(alias="datasetId")] = ""
    r"""The ID of the dataset to train the element on."""

    description: OptionalNullable[str] = ""
    r"""The description of the element."""

    instance_prompt: Optional[str] = "a character"
    r"""The instance prompt to use during training.Try “a” by a noun. E.g. a castle"""

    learning_rate: Optional[float] = 0.000001
    r"""The speed of element learns."""

    lora_focus: Optional[str] = "General"
    r"""The category determines how the element will be trained. Options are 'General' | 'Character' | 'Style' | 'Object'."""

    name: Optional[str] = "placeholder"
    r"""The name of the element."""

    num_train_epochs: Optional[int] = 100
    r"""The number of times the entire training dataset is passed through the element."""

    resolution: OptionalNullable[int] = 1024
    r"""The resolution for training. Must be 1024."""

    sd_version: Optional[SdVersions] = SdVersions.SDXL_0_9
    r"""The base version of stable diffusion to use if not using a custom model."""

    train_text_encoder: Optional[bool] = True
    r"""Whether or not encode the train text."""

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = [
            "datasetId",
            "description",
            "instance_prompt",
            "learning_rate",
            "lora_focus",
            "name",
            "num_train_epochs",
            "resolution",
            "sd_version",
            "train_text_encoder",
        ]
        nullable_fields = ["description", "resolution"]
        null_default_fields = []

        serialized = handler(self)

        m = {}

        for n, f in self.model_fields.items():
            k = f.alias or n
            val = serialized.get(k)
            serialized.pop(k, None)

            optional_nullable = k in optional_fields and k in nullable_fields
            is_set = (
                self.__pydantic_fields_set__.intersection({n})
                or k in null_default_fields
            )  # pylint: disable=no-member

            if val is not None and val != UNSET_SENTINEL:
                m[k] = val
            elif val != UNSET_SENTINEL and (
                not k in optional_fields or (optional_nullable and is_set)
            ):
                m[k] = val

        return m


class SDTrainingOutputTypedDict(TypedDict):
    api_credit_cost: NotRequired[Nullable[int]]
    r"""API Credits Cost for Model Training. Available for Production API Users."""
    user_lora_id: NotRequired[int]


class SDTrainingOutput(BaseModel):
    api_credit_cost: Annotated[
        OptionalNullable[int], pydantic.Field(alias="apiCreditCost")
    ] = UNSET
    r"""API Credits Cost for Model Training. Available for Production API Users."""

    user_lora_id: Annotated[Optional[int], pydantic.Field(alias="userLoraId")] = None

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = ["apiCreditCost", "userLoraId"]
        nullable_fields = ["apiCreditCost"]
        null_default_fields = []

        serialized = handler(self)

        m = {}

        for n, f in self.model_fields.items():
            k = f.alias or n
            val = serialized.get(k)
            serialized.pop(k, None)

            optional_nullable = k in optional_fields and k in nullable_fields
            is_set = (
                self.__pydantic_fields_set__.intersection({n})
                or k in null_default_fields
            )  # pylint: disable=no-member

            if val is not None and val != UNSET_SENTINEL:
                m[k] = val
            elif val != UNSET_SENTINEL and (
                not k in optional_fields or (optional_nullable and is_set)
            ):
                m[k] = val

        return m


class CreateElementResponseBodyTypedDict(TypedDict):
    r"""Responses for POST /elements."""

    sd_training_job: NotRequired[Nullable[SDTrainingOutputTypedDict]]


class CreateElementResponseBody(BaseModel):
    r"""Responses for POST /elements."""

    sd_training_job: Annotated[
        OptionalNullable[SDTrainingOutput], pydantic.Field(alias="sdTrainingJob")
    ] = UNSET

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = ["sdTrainingJob"]
        nullable_fields = ["sdTrainingJob"]
        null_default_fields = []

        serialized = handler(self)

        m = {}

        for n, f in self.model_fields.items():
            k = f.alias or n
            val = serialized.get(k)
            serialized.pop(k, None)

            optional_nullable = k in optional_fields and k in nullable_fields
            is_set = (
                self.__pydantic_fields_set__.intersection({n})
                or k in null_default_fields
            )  # pylint: disable=no-member

            if val is not None and val != UNSET_SENTINEL:
                m[k] = val
            elif val != UNSET_SENTINEL and (
                not k in optional_fields or (optional_nullable and is_set)
            ):
                m[k] = val

        return m


class CreateElementResponseTypedDict(TypedDict):
    content_type: str
    r"""HTTP response content type for this operation"""
    status_code: int
    r"""HTTP response status code for this operation"""
    raw_response: httpx.Response
    r"""Raw HTTP response; suitable for custom response parsing"""
    object: NotRequired[CreateElementResponseBodyTypedDict]
    r"""Responses for POST /elements."""


class CreateElementResponse(BaseModel):
    content_type: str
    r"""HTTP response content type for this operation"""

    status_code: int
    r"""HTTP response status code for this operation"""

    raw_response: httpx.Response
    r"""Raw HTTP response; suitable for custom response parsing"""

    object: Optional[CreateElementResponseBody] = None
    r"""Responses for POST /elements."""
