from PIL import Image
from pydantic import BaseModel, field_validator

from screensuite.benchmarks.singlestep.common_models import Operation


class Sample(BaseModel):
    """Represents a sample from the dataset."""

    operation: Operation
    screenshot: Image.Image
    coordinates: tuple[int, int, int, int]
    image_size: tuple[int, int]

    @field_validator("screenshot", mode="before")
    def validate_screenshot(cls, v):
        if v is None:
            return b"\x00"
        elif isinstance(v, bytes):
            return v
        else:
            return v["bytes"]
