from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from pipelex.cogt.extract.bounding_box import BoundingBox


class GatewayExtractRequestParams(BaseModel):
    should_include_images: bool


class GatewayExtractImageAzure(BaseModel):
    """An extracted image with its base64 content and bounding box."""

    base64_str: str = Field(description="Base64-encoded image content")
    mime_type: str = Field(description="MIME type of the image")
    bounding_box: BoundingBox | None = Field(description="Bounding box coordinates of the image in the page", default=None)
    caption: str | None = Field(description="Caption text associated with the image", default=None)


class GatewayExtractPageAzure(BaseModel):
    """Result for a single page extracted by Azure Doc Intel."""

    index: int = Field(description="The index of the page")
    markdown: str = Field(description="The markdown content of the page")
    images: list[GatewayExtractImageAzure] = Field(description="The images on the page")


class ImageDimensions(BaseModel):
    """Image dimensions and file size."""

    width: int = Field(description="Image width in pixels")
    height: int = Field(description="Image height in pixels")
    bytes: int = Field(description="Image file size in bytes")


class ImageInfo(BaseModel):
    """Information about the processed image."""

    original: ImageDimensions = Field(description="Original image dimensions before processing")
    processed: ImageDimensions = Field(description="Processed image dimensions (after scaling if applied)")
    scaled_down: bool = Field(description="Whether the image was scaled down to fit limits")


class GatewayExtractPageDeepseek(BaseModel):
    """Result for a single page extracted by DeepSeek-OCR."""

    index: int = Field(description="The index of the page (0-based)")
    markdown: str = Field(description="The markdown content of the page")
    source_image_info: ImageInfo | None = Field(default=None, description="Information about the source image for this page")


class GatewayExtractImageMistral(BaseModel):
    """An extracted image with its base64 content and bounding box."""

    id: str | None = Field(description="The ID of the image", default=None)
    top_left_x: int | None = Field(description="The X coordinate of the top left corner of the image", default=None)
    top_left_y: int | None = Field(description="The Y coordinate of the top left corner of the image", default=None)
    bottom_right_x: int | None = Field(description="The X coordinate of the bottom right corner of the image", default=None)
    bottom_right_y: int | None = Field(description="The Y coordinate of the bottom right corner of the image", default=None)
    image_base64: str | None = Field(description="Base64-encoded image content", default=None)
    image_annotation: str | None = Field(description="Annotation text associated with the image", default=None)


class GatewayExtractPageMistral(BaseModel):
    """Result for a single page extracted by Mistral Doc AI."""

    index: int = Field(description="The index of the page")
    markdown: str = Field(description="The markdown content of the page")
    images: list[GatewayExtractImageMistral] = Field(description="The images on the page")


class GatewayImgGenAzureGptImage(BaseModel):
    """Result for a single image generated by Azure GPT Image."""

    model_config = ConfigDict(extra="ignore")

    size: str = Field(description="Size of the image")
    output_format: str = Field(description="Output format of the image")


class GatewayImgGenAzureFlux2Pro(BaseModel):
    """Result for a single image generated by Azure Flux 2 Pro."""

    model_config = ConfigDict(extra="ignore")

    request_meta: dict[str, Any] = Field(description="Request metadata", default_factory=dict)
