"""Pydantic models for API request bodies."""

from typing import Any, Literal

from pydantic import BaseModel


class TransformData(BaseModel):
    """Data for transform execution.

    This model accepts arbitrary additional fields beyond 'image'.
    For example, you can pass custom fields like 'test_bbox' for transforms
    that require additional data (e.g., RandomCropNearBBox).

    Example:
        {
            "image": "base64_encoded_image",
            "test_bbox": [0, 5, 10, 20],  # Custom field for RandomCropNearBBox
            "any_other_field": "value"     # Any other data needed by transforms
        }
    """

    image: str | None = None  # base64 encoded image

    # Allow arbitrary additional fields
    model_config = {
        "extra": "allow",
    }


class TransformRequest(BaseModel):
    """Request for transform execution."""

    transform: str
    params: dict[str, Any]
    data: TransformData = TransformData()  # Contains image and potentially other data in the future
    bbox_params: dict[str, Any] | None = None
    keypoint_params: dict[str, Any] | None = None
    targets: list[Literal["image", "mask", "bboxes", "keypoints"]]
