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

from __future__ import annotations
from enum import Enum
from lambdadb.types import BaseModel, UNSET_SENTINEL
from lambdadb.utils import get_discriminator
import pydantic
from pydantic import Discriminator, Tag, model_serializer
from typing import Any, Dict, List, Optional, Union
from typing_extensions import Annotated, NotRequired, TypeAliasType, TypedDict


class TypeObject(str, Enum):
    OBJECT = "object"


class IndexConfigsObjectTypedDict(TypedDict):
    type: TypeObject
    object_index_configs: Dict[str, Any]


class IndexConfigsObject(BaseModel):
    type: TypeObject

    object_index_configs: Annotated[
        Dict[str, Any], pydantic.Field(alias="objectIndexConfigs")
    ]


class Type(str, Enum):
    KEYWORD = "keyword"
    LONG = "long"
    DOUBLE = "double"
    DATETIME = "datetime"
    BOOLEAN = "boolean"
    SPARSE_VECTOR = "sparseVector"


class IndexConfigsTypedDict(TypedDict):
    r"""Types that do not need additional parameters."""

    type: Type


class IndexConfigs(BaseModel):
    r"""Types that do not need additional parameters."""

    type: Type


class TypeVector(str, Enum):
    VECTOR = "vector"


class Similarity(str, Enum):
    r"""Vector similarity metric."""

    COSINE = "cosine"
    EUCLIDEAN = "euclidean"
    DOT_PRODUCT = "dot_product"
    MAX_INNER_PRODUCT = "max_inner_product"


class IndexConfigsVectorTypedDict(TypedDict):
    type: TypeVector
    dimensions: int
    r"""Vector dimensions."""
    similarity: NotRequired[Similarity]
    r"""Vector similarity metric."""


class IndexConfigsVector(BaseModel):
    type: TypeVector

    dimensions: int
    r"""Vector dimensions."""

    similarity: Optional[Similarity] = Similarity.COSINE
    r"""Vector similarity metric."""

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = set(["similarity"])
        serialized = handler(self)
        m = {}

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

            if val != UNSET_SENTINEL:
                if val is not None or k not in optional_fields:
                    m[k] = val

        return m


class TypeText(str, Enum):
    TEXT = "text"


class Analyzer(str, Enum):
    STANDARD = "standard"
    KOREAN = "korean"
    JAPANESE = "japanese"
    ENGLISH = "english"


class IndexConfigsTextTypedDict(TypedDict):
    type: TypeText
    analyzers: NotRequired[List[Analyzer]]
    r"""Analyzers."""


class IndexConfigsText(BaseModel):
    type: TypeText

    analyzers: Optional[List[Analyzer]] = None
    r"""Analyzers."""

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = set(["analyzers"])
        serialized = handler(self)
        m = {}

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

            if val != UNSET_SENTINEL:
                if val is not None or k not in optional_fields:
                    m[k] = val

        return m


IndexConfigsUnionTypedDict = TypeAliasType(
    "IndexConfigsUnionTypedDict",
    Union[
        IndexConfigsTypedDict,
        IndexConfigsTextTypedDict,
        IndexConfigsObjectTypedDict,
        IndexConfigsVectorTypedDict,
    ],
)


IndexConfigsUnion = Annotated[
    Union[
        Annotated[IndexConfigsText, Tag("text")],
        Annotated[IndexConfigsVector, Tag("vector")],
        Annotated[IndexConfigs, Tag("keyword")],
        Annotated[IndexConfigs, Tag("long")],
        Annotated[IndexConfigs, Tag("double")],
        Annotated[IndexConfigs, Tag("datetime")],
        Annotated[IndexConfigs, Tag("boolean")],
        Annotated[IndexConfigs, Tag("sparseVector")],
        Annotated[IndexConfigsObject, Tag("object")],
    ],
    Discriminator(lambda m: get_discriminator(m, "type", "type")),
]
