"""
Pydantic models for request/response validation.
"""
import enum
from typing import Optional

import pydantic


class ImageData(pydantic.BaseModel):
    """Image data model for base64 encoded images."""
    mediaType: str
    data: str


class ImageRequest(pydantic.BaseModel):
    """Request model for image classification."""
    image: ImageData


class Labels(enum.IntEnum):
    Natural = 0
    FullySynthesized = 1
    LocallyEdited = 2
    LocallySynthesized = 3


class LocalizationMask(pydantic.BaseModel):
    """A bit mask indicating which pixels are manipulated / synthesized.

    A bit value of ``1`` means that the model believes the corresponding pixel
    has been edited or synthesized (i.e., its label would be non-zero).
    A bit value of ``0`` means that the model believes the pixel is unaltered.

    The mask ``.width`` and ``.height`` should be the same as the input image.
    Extra bits at the end of ``.bitsRowMajor`` after the first
    ``width * height`` bits are **ignored**; for simplicity/efficiency,
    you should encode your bit mask into a byte array and not worry if the
    final byte isn't "full", then convert the byte array to base64.
    """

    width: int = pydantic.Field(
        description="The width of the mask."
    )

    height: int = pydantic.Field(
        description="The height of the mask."
    )
    
    bitsRowMajor: str = pydantic.Field(
        description="A base64 string encoding the bit mask in row-major order.",
        # Canonical base64 encoding
        # https://stackoverflow.com/a/64467300/3709935
        pattern=r"^(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/][AQgw]==|[A-Za-z0-9+/]{2}[AEIMQUYcgkosw048]=)?$",
    )


class PredictionResponse(pydantic.BaseModel):
    """Response model for synthetic image classification results.
    
    Detector models will be scored primarily on their ability to classify the
    entire image into 1 of the 4 label categories::

        0: (Natural) The image is natural / unaltered.
        1: (FullySynthesized) The entire image was synthesized by e.g., a
            generative image model.
        2: (LocallyEdited) The image is a natural image where a portion has
            been edited using traditional photo editing techniques such as
            splicing.
        3: (LocallySynthesized) The image is a natural image where a portion
            has been replaced by synthesized content.
    """

    logprobs: list[float] = pydantic.Field(
        description="The log-probabilities for each of the 4 possible labels.",
        min_length=4,
        max_length=4,
    )

    localizationMask: Optional[LocalizationMask] = pydantic.Field(
        description="A bit mask localizing predicted edits. Models that are"
        " not capable of localization may omit this field. It may also be"
        " omitted if the predicted label is ``0`` or ``1``, in which case the"
        " mask will be assumed to be all 0's or all 1's, as appropriate."
    )