"""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.types import BaseModel, Nullable, OptionalNullable, UNSET, UNSET_SENTINEL
from pydantic import model_serializer
from typing import List, 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"""


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


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

    @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",
        ]
        nullable_fields = [
            "model",
            "pages",
            "include_image_base64",
            "image_limit",
            "image_min_size",
            "bbox_annotation_format",
            "document_annotation_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