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

from __future__ import annotations
from .anthropiccachecontroldirective import (
    AnthropicCacheControlDirective,
    AnthropicCacheControlDirectiveTypedDict,
)
from .anthropiccitationcharlocationparam import (
    AnthropicCitationCharLocationParam,
    AnthropicCitationCharLocationParamTypedDict,
)
from .anthropiccitationcontentblocklocationparam import (
    AnthropicCitationContentBlockLocationParam,
    AnthropicCitationContentBlockLocationParamTypedDict,
)
from .anthropiccitationpagelocationparam import (
    AnthropicCitationPageLocationParam,
    AnthropicCitationPageLocationParamTypedDict,
)
from .anthropiccitationsearchresultlocation import (
    AnthropicCitationSearchResultLocation,
    AnthropicCitationSearchResultLocationTypedDict,
)
from .anthropiccitationwebsearchresultlocation import (
    AnthropicCitationWebSearchResultLocation,
    AnthropicCitationWebSearchResultLocationTypedDict,
)
from openrouter.types import (
    BaseModel,
    Nullable,
    OptionalNullable,
    UNSET,
    UNSET_SENTINEL,
)
from openrouter.utils import get_discriminator
from pydantic import Discriminator, Tag, model_serializer
from typing import List, Literal, Optional, Union
from typing_extensions import Annotated, NotRequired, TypeAliasType, TypedDict


CitationTypedDict = TypeAliasType(
    "CitationTypedDict",
    Union[
        AnthropicCitationWebSearchResultLocationTypedDict,
        AnthropicCitationCharLocationParamTypedDict,
        AnthropicCitationPageLocationParamTypedDict,
        AnthropicCitationContentBlockLocationParamTypedDict,
        AnthropicCitationSearchResultLocationTypedDict,
    ],
)


Citation = Annotated[
    Union[
        Annotated[AnthropicCitationCharLocationParam, Tag("char_location")],
        Annotated[
            AnthropicCitationContentBlockLocationParam, Tag("content_block_location")
        ],
        Annotated[AnthropicCitationPageLocationParam, Tag("page_location")],
        Annotated[AnthropicCitationSearchResultLocation, Tag("search_result_location")],
        Annotated[
            AnthropicCitationWebSearchResultLocation, Tag("web_search_result_location")
        ],
    ],
    Discriminator(lambda m: get_discriminator(m, "type", "type")),
]


AnthropicTextBlockParamType = Literal["text",]


class AnthropicTextBlockParamTypedDict(TypedDict):
    text: str
    type: AnthropicTextBlockParamType
    cache_control: NotRequired[AnthropicCacheControlDirectiveTypedDict]
    r"""Enable automatic prompt caching. When set at the top level, the system automatically applies cache breakpoints to the last cacheable block in the request. Currently supported for Anthropic Claude models."""
    citations: NotRequired[Nullable[List[CitationTypedDict]]]


class AnthropicTextBlockParam(BaseModel):
    text: str

    type: AnthropicTextBlockParamType

    cache_control: Optional[AnthropicCacheControlDirective] = None
    r"""Enable automatic prompt caching. When set at the top level, the system automatically applies cache breakpoints to the last cacheable block in the request. Currently supported for Anthropic Claude models."""

    citations: OptionalNullable[List[Citation]] = UNSET

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = ["cache_control", "citations"]
        nullable_fields = ["citations"]
        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
