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

from __future__ import annotations
from .responseformatjsonschema import (
    ResponseFormatJSONSchema,
    ResponseFormatJSONSchemaTypedDict,
)
from .responseformattextgrammar import (
    ResponseFormatTextGrammar,
    ResponseFormatTextGrammarTypedDict,
)
from openrouter.types import (
    BaseModel,
    Nullable,
    OptionalNullable,
    UNSET,
    UNSET_SENTINEL,
)
from openrouter.utils import get_discriminator, validate_const
import pydantic
from pydantic import Discriminator, Tag, model_serializer
from pydantic.functional_validators import AfterValidator
from typing import Dict, List, Literal, Optional, Union
from typing_extensions import Annotated, NotRequired, TypeAliasType, TypedDict


PromptTypedDict = TypeAliasType(
    "PromptTypedDict", Union[str, List[str], List[float], List[List[float]]]
)


Prompt = TypeAliasType("Prompt", Union[str, List[str], List[float], List[List[float]]])


CompletionCreateParamsStopTypedDict = TypeAliasType(
    "CompletionCreateParamsStopTypedDict", Union[str, List[str]]
)


CompletionCreateParamsStop = TypeAliasType(
    "CompletionCreateParamsStop", Union[str, List[str]]
)


class StreamOptionsTypedDict(TypedDict):
    include_usage: NotRequired[Nullable[bool]]


class StreamOptions(BaseModel):
    include_usage: OptionalNullable[bool] = UNSET

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

        serialized = handler(self)

        m = {}

        for n, f in type(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 CompletionCreateParamsResponseFormatPythonTypedDict(TypedDict):
    type: Literal["python"]


class CompletionCreateParamsResponseFormatPython(BaseModel):
    TYPE: Annotated[
        Annotated[Literal["python"], AfterValidator(validate_const("python"))],
        pydantic.Field(alias="type"),
    ] = "python"


class CompletionCreateParamsResponseFormatJSONObjectTypedDict(TypedDict):
    type: Literal["json_object"]


class CompletionCreateParamsResponseFormatJSONObject(BaseModel):
    TYPE: Annotated[
        Annotated[
            Literal["json_object"], AfterValidator(validate_const("json_object"))
        ],
        pydantic.Field(alias="type"),
    ] = "json_object"


class CompletionCreateParamsResponseFormatTextTypedDict(TypedDict):
    type: Literal["text"]


class CompletionCreateParamsResponseFormatText(BaseModel):
    TYPE: Annotated[
        Annotated[Literal["text"], AfterValidator(validate_const("text"))],
        pydantic.Field(alias="type"),
    ] = "text"


CompletionCreateParamsResponseFormatUnionTypedDict = TypeAliasType(
    "CompletionCreateParamsResponseFormatUnionTypedDict",
    Union[
        CompletionCreateParamsResponseFormatTextTypedDict,
        CompletionCreateParamsResponseFormatJSONObjectTypedDict,
        CompletionCreateParamsResponseFormatPythonTypedDict,
        ResponseFormatJSONSchemaTypedDict,
        ResponseFormatTextGrammarTypedDict,
    ],
)


CompletionCreateParamsResponseFormatUnion = Annotated[
    Union[
        Annotated[CompletionCreateParamsResponseFormatText, Tag("text")],
        Annotated[CompletionCreateParamsResponseFormatJSONObject, Tag("json_object")],
        Annotated[ResponseFormatJSONSchema, Tag("json_schema")],
        Annotated[ResponseFormatTextGrammar, Tag("grammar")],
        Annotated[CompletionCreateParamsResponseFormatPython, Tag("python")],
    ],
    Discriminator(lambda m: get_discriminator(m, "type", "type")),
]


class CompletionCreateParamsTypedDict(TypedDict):
    prompt: PromptTypedDict
    model: NotRequired[str]
    models: NotRequired[List[str]]
    best_of: NotRequired[Nullable[int]]
    echo: NotRequired[Nullable[bool]]
    frequency_penalty: NotRequired[Nullable[float]]
    logit_bias: NotRequired[Nullable[Dict[str, float]]]
    logprobs: NotRequired[Nullable[int]]
    max_tokens: NotRequired[Nullable[int]]
    n: NotRequired[Nullable[int]]
    presence_penalty: NotRequired[Nullable[float]]
    seed: NotRequired[Nullable[int]]
    stop: NotRequired[Nullable[CompletionCreateParamsStopTypedDict]]
    stream: NotRequired[bool]
    stream_options: NotRequired[Nullable[StreamOptionsTypedDict]]
    suffix: NotRequired[Nullable[str]]
    temperature: NotRequired[Nullable[float]]
    top_p: NotRequired[Nullable[float]]
    user: NotRequired[str]
    metadata: NotRequired[Nullable[Dict[str, str]]]
    response_format: NotRequired[
        Nullable[CompletionCreateParamsResponseFormatUnionTypedDict]
    ]


class CompletionCreateParams(BaseModel):
    prompt: Prompt

    model: Optional[str] = None

    models: Optional[List[str]] = None

    best_of: OptionalNullable[int] = UNSET

    echo: OptionalNullable[bool] = UNSET

    frequency_penalty: OptionalNullable[float] = UNSET

    logit_bias: OptionalNullable[Dict[str, float]] = UNSET

    logprobs: OptionalNullable[int] = UNSET

    max_tokens: OptionalNullable[int] = UNSET

    n: OptionalNullable[int] = UNSET

    presence_penalty: OptionalNullable[float] = UNSET

    seed: OptionalNullable[int] = UNSET

    stop: OptionalNullable[CompletionCreateParamsStop] = UNSET

    stream: Optional[bool] = False

    stream_options: OptionalNullable[StreamOptions] = UNSET

    suffix: OptionalNullable[str] = UNSET

    temperature: OptionalNullable[float] = UNSET

    top_p: OptionalNullable[float] = UNSET

    user: Optional[str] = None

    metadata: OptionalNullable[Dict[str, str]] = UNSET

    response_format: OptionalNullable[CompletionCreateParamsResponseFormatUnion] = UNSET

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = [
            "model",
            "models",
            "best_of",
            "echo",
            "frequency_penalty",
            "logit_bias",
            "logprobs",
            "max_tokens",
            "n",
            "presence_penalty",
            "seed",
            "stop",
            "stream",
            "stream_options",
            "suffix",
            "temperature",
            "top_p",
            "user",
            "metadata",
            "response_format",
        ]
        nullable_fields = [
            "best_of",
            "echo",
            "frequency_penalty",
            "logit_bias",
            "logprobs",
            "max_tokens",
            "n",
            "presence_penalty",
            "seed",
            "stop",
            "stream_options",
            "suffix",
            "temperature",
            "top_p",
            "metadata",
            "response_format",
        ]
        null_default_fields = []

        serialized = handler(self)

        m = {}

        for n, f in type(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