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

from __future__ import annotations
from .documenturlchunk import DocumentURLChunk, DocumentURLChunkTypedDict
from .filechunk import FileChunk, FileChunkTypedDict
from .imageurlchunk import ImageURLChunk, ImageURLChunkTypedDict
from .responseformat import ResponseFormat, ResponseFormatTypedDict
from mistralai.azure.client.types import (
    BaseModel,
    Nullable,
    OptionalNullable,
    UNSET,
    UNSET_SENTINEL,
)
from pydantic import model_serializer
from typing import List, Literal, Optional, Union
from typing_extensions import NotRequired, TypeAliasType, TypedDict


DocumentTypedDict = TypeAliasType(
    "DocumentTypedDict",
    Union[FileChunkTypedDict, ImageURLChunkTypedDict, DocumentURLChunkTypedDict],
)
r"""Document to run OCR on"""


Document = TypeAliasType("Document", Union[FileChunk, ImageURLChunk, DocumentURLChunk])
r"""Document to run OCR on"""


TableFormat = Literal[
    "markdown",
    "html",
]


class OCRRequestTypedDict(TypedDict):
    model: Nullable[str]
    document: DocumentTypedDict
    r"""Document to run OCR on"""
    id: NotRequired[str]
    pages: NotRequired[Nullable[List[int]]]
    r"""Specific pages user wants to process in various formats: single number, range, or list of both. Starts from 0"""
    include_image_base64: NotRequired[Nullable[bool]]
    r"""Include image URLs in response"""
    image_limit: NotRequired[Nullable[int]]
    r"""Max images to extract"""
    image_min_size: NotRequired[Nullable[int]]
    r"""Minimum height and width of image to extract"""
    bbox_annotation_format: NotRequired[Nullable[ResponseFormatTypedDict]]
    r"""Structured output class for extracting useful information from each extracted bounding box / image from document. Only json_schema is valid for this field"""
    document_annotation_format: NotRequired[Nullable[ResponseFormatTypedDict]]
    r"""Structured output class for extracting useful information from the entire document. Only json_schema is valid for this field"""
    document_annotation_prompt: NotRequired[Nullable[str]]
    r"""Optional prompt to guide the model in extracting structured output from the entire document. A document_annotation_format must be provided."""
    table_format: NotRequired[Nullable[TableFormat]]
    extract_header: NotRequired[bool]
    extract_footer: NotRequired[bool]


class OCRRequest(BaseModel):
    model: Nullable[str]

    document: Document
    r"""Document to run OCR on"""

    id: Optional[str] = None

    pages: OptionalNullable[List[int]] = UNSET
    r"""Specific pages user wants to process in various formats: single number, range, or list of both. Starts from 0"""

    include_image_base64: OptionalNullable[bool] = UNSET
    r"""Include image URLs in response"""

    image_limit: OptionalNullable[int] = UNSET
    r"""Max images to extract"""

    image_min_size: OptionalNullable[int] = UNSET
    r"""Minimum height and width of image to extract"""

    bbox_annotation_format: OptionalNullable[ResponseFormat] = UNSET
    r"""Structured output class for extracting useful information from each extracted bounding box / image from document. Only json_schema is valid for this field"""

    document_annotation_format: OptionalNullable[ResponseFormat] = UNSET
    r"""Structured output class for extracting useful information from the entire document. Only json_schema is valid for this field"""

    document_annotation_prompt: OptionalNullable[str] = UNSET
    r"""Optional prompt to guide the model in extracting structured output from the entire document. A document_annotation_format must be provided."""

    table_format: OptionalNullable[TableFormat] = UNSET

    extract_header: Optional[bool] = None

    extract_footer: Optional[bool] = None

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = set(
            [
                "id",
                "pages",
                "include_image_base64",
                "image_limit",
                "image_min_size",
                "bbox_annotation_format",
                "document_annotation_format",
                "document_annotation_prompt",
                "table_format",
                "extract_header",
                "extract_footer",
            ]
        )
        nullable_fields = set(
            [
                "model",
                "pages",
                "include_image_base64",
                "image_limit",
                "image_min_size",
                "bbox_annotation_format",
                "document_annotation_format",
                "document_annotation_prompt",
                "table_format",
            ]
        )
        serialized = handler(self)
        m = {}

        for n, f in type(self).model_fields.items():
            k = f.alias or n
            val = serialized.get(k)
            is_nullable_and_explicitly_set = (
                k in nullable_fields
                and (self.__pydantic_fields_set__.intersection({n}))  # pylint: disable=no-member
            )

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

        return m