"""Pydantic models related to image data and annotations."""

from typing import Literal

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel


class BBoxData(BaseModel):
    """Represents bounding box data."""

    x_min: float
    y_min: float
    x_max: float
    y_max: float
    label: str | None = None


class KeypointData(BaseModel):
    """Represents keypoint data."""

    x: float
    y: float
    label: str | None = None


# Add Annotations class here
class Annotations:
    """Container for all annotation types."""

    def __init__(
        self,
        bboxes: list[BBoxData],
        keypoints: list[KeypointData],
        mask: npt.NDArray[np.uint8] | None = None,
        image_width: int = 0,
        image_height: int = 0,
    ) -> None:
        self.bboxes = bboxes
        self.keypoints = keypoints
        self.mask = mask
        self.image_width = image_width
        self.image_height = image_height

    class Config:
        arbitrary_types_allowed = True  # Needed for npt.NDArray


Target = Literal["image", "mask", "bboxes", "keypoints"]


class ImageData(BaseModel):
    """Unified model for primary or reference image data with visualizations."""

    # Common fields
    image: str  # base64 encoded image
    mask: str | None = None  # base64 encoded image
    bboxes: list[BBoxData] | None = None
    keypoints: list[KeypointData] | None = None


class ImageSet(BaseModel):
    """Configuration for an image set and its optional elements."""

    name: str  # Identifier for the image set
    image_filename: str  # Filename of the image
    annotation_filename: str | None = None  # Filename of the annotation file (if any)

    model_config = {
        "frozen": True,  # Make instances hashable for caching
    }
