"""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.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"""
    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"""

    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 = [
            "id",
            "pages",
            "include_image_base64",
            "image_limit",
            "image_min_size",
            "bbox_annotation_format",
            "document_annotation_format",
            "table_format",
            "extract_header",
            "extract_footer",
        ]
        nullable_fields = [
            "model",
            "pages",
            "include_image_base64",
            "image_limit",
            "image_min_size",
            "bbox_annotation_format",
            "document_annotation_format",
            "table_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
