src/editable_image/models.py

from __future__ import annotations

from enum import StrEnum
from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator


Pixel = Annotated[float, Field(ge=0)]
DEFAULT_SNAP_ANGLES = (0.0, 45.0, 90.0, -45.0, -90.0)


class CleanupMode(StrEnum):
    OPENCV = "opencv"
    LAMA = "lama"


class Device(StrEnum):
    AUTO = "auto"
    CPU = "cpu"
    CUDA = "cuda"
    MPS = "mps"


class BitmapExtractor(StrEnum):
    OPENCV = "opencv"
    SAM2 = "sam2"


class Point(BaseModel):
    x: Pixel
    y: Pixel


class Box(BaseModel):
    x: Pixel
    y: Pixel
    width: Annotated[float, Field(gt=0)]
    height: Annotated[float, Field(gt=0)]


class TextStyle(BaseModel):
    font_family: str = "Lato"
    font_size: Annotated[float, Field(ge=4, le=1000)] = 16
    font_weight: Literal[400, 700] = 400
    italic: bool = False
    fill: str = "#111111"
    opacity: Annotated[float, Field(ge=0, le=1)] = 1
    stroke: str = "none"
    stroke_width: Annotated[float, Field(ge=0, le=50)] = 0
    letter_spacing: Annotated[float, Field(ge=-20, le=100)] = 0
    anchor: Literal["start", "middle", "end"] = "start"

    @field_validator("fill", "stroke")
    @classmethod
    def validate_color(cls, value: str) -> str:
        if value == "none" or (
            len(value) in {4, 7, 9} and value.startswith("#")
        ):
            return value
        raise ValueError("expected a CSS hexadecimal color or 'none'")


class TextRegion(BaseModel):
    model_config = ConfigDict(extra="forbid")

    id: str
    text: str
    confidence: Annotated[float, Field(ge=0, le=1)]
    source_quad: Annotated[list[Point], Field(min_length=4, max_length=4)]
    overlay_box: Box
    rotation: Annotated[float, Field(ge=-180, le=180)] = 0
    remove_text: bool = True
    include_overlay: bool = True
    mask_padding: Annotated[float, Field(ge=0, le=100)] = 3
    style: TextStyle = Field(default_factory=TextStyle)
    z_index: int = 1000


class SolidPaint(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal["solid"] = "solid"
    color: str

    @field_validator("color")
    @classmethod
    def validate_color(cls, value: str) -> str:
        if len(value) in {4, 7, 9} and value.startswith("#"):
            return value
        raise ValueError("expected a CSS hexadecimal color")


class GradientStop(BaseModel):
    offset: Annotated[float, Field(ge=0, le=1)]
    color: str
    opacity: Annotated[float, Field(ge=0, le=1)] = 1

    @field_validator("color")
    @classmethod
    def validate_color(cls, value: str) -> str:
        if len(value) in {4, 7, 9} and value.startswith("#"):
            return value
        raise ValueError("expected a CSS hexadecimal color")


class LinearGradientPaint(BaseModel):
    model_config = ConfigDict(extra="forbid")

    type: Literal["linear-gradient"] = "linear-gradient"
    angle: Annotated[float, Field(ge=-180, le=180)] = 0
    stops: Annotated[list[GradientStop], Field(min_length=2, max_length=8)]


Paint = SolidPaint | LinearGradientPaint


class ShapeStyle(BaseModel):
    fill: Paint = Field(default_factory=lambda: SolidPaint(color="#000000"))
    fill_opacity: Annotated[float, Field(ge=0, le=1)] = 1
    stroke: str = "none"
    stroke_width: Annotated[float, Field(ge=0, le=100)] = 0
    stroke_opacity: Annotated[float, Field(ge=0, le=1)] = 1
    linecap: Literal["butt", "round", "square"] = "round"
    linejoin: Literal["miter", "round", "bevel"] = "round"

    @field_validator("stroke")
    @classmethod
    def validate_stroke(cls, value: str) -> str:
        if value == "none" or (len(value) in {4, 7, 9} and value.startswith("#")):
            return value
        raise ValueError("expected a CSS hexadecimal color or 'none'")


class RectGeometry(BaseModel):
    type: Literal["rect"] = "rect"
    box: Box
    rx: Annotated[float, Field(ge=0)] = 0
    ry: Annotated[float, Field(ge=0)] = 0
    rotation: Annotated[float, Field(ge=-180, le=180)] = 0


class EllipseGeometry(BaseModel):
    type: Literal["ellipse"] = "ellipse"
    cx: Pixel
    cy: Pixel
    rx: Annotated[float, Field(gt=0)]
    ry: Annotated[float, Field(gt=0)]
    rotation: Annotated[float, Field(ge=-180, le=180)] = 0


class PolygonGeometry(BaseModel):
    type: Literal["polygon"] = "polygon"
    points: Annotated[list[Point], Field(min_length=3)]


class PathGeometry(BaseModel):
    type: Literal["path"] = "path"
    points: Annotated[list[Point], Field(min_length=2)]
    closed: bool = False
    smooth: bool = False
    arrow_start: bool = False
    arrow_end: bool = False


ShapeGeometry = RectGeometry | EllipseGeometry | PolygonGeometry | PathGeometry


class ShapeRegion(BaseModel):
    model_config = ConfigDict(extra="forbid")

    id: str
    confidence: Annotated[float, Field(ge=0, le=1)]
    geometry: ShapeGeometry
    source_contours: list[list[Point]] = Field(default_factory=list)
    style: ShapeStyle = Field(default_factory=ShapeStyle)
    remove_source: bool = True
    include_overlay: bool = True
    mask_padding: Annotated[float, Field(ge=0, le=100)] = 1
    z_index: int = 100
    is_container: bool = False
    semantic_role: Literal["shape", "arrow"] = "shape"


class BitmapRegion(BaseModel):
    model_config = ConfigDict(extra="forbid")

    id: str
    asset_name: str
    box: Box
    opacity: Annotated[float, Field(ge=0, le=1)] = 1
    z_index: int = 500


class MaskStroke(BaseModel):
    operation: Literal["add", "erase"]
    radius: Annotated[float, Field(gt=0, le=500)]
    points: Annotated[list[Point], Field(min_length=1)]


class Canvas(BaseModel):
    width: Annotated[int, Field(gt=0)]
    height: Annotated[int, Field(gt=0)]


class DocumentSpec(BaseModel):
    model_config = ConfigDict(extra="forbid")

    version: Literal[1, 2] = 1
    source_name: str
    canvas: Canvas
    regions: list[TextRegion]
    shapes: list[ShapeRegion] = Field(default_factory=list)
    bitmap_regions: list[BitmapRegion] = Field(default_factory=list)
    mask_strokes: list[MaskStroke] = Field(default_factory=list)


class ProcessOptions(BaseModel):
    confidence: Annotated[float, Field(ge=0, le=1)] = 0.5
    device: Device = Device.AUTO
    cleanup: CleanupMode = CleanupMode.OPENCV
    max_pixels: Annotated[int, Field(gt=0)] = 40_000_000
    snap_angles: Annotated[
        list[Annotated[float, Field(ge=-180, le=180)]],
        Field(min_length=1),
    ] = Field(default_factory=lambda: list(DEFAULT_SNAP_ANGLES))
    snap_tolerance: Annotated[float, Field(ge=0, le=180)] = 6.0
    snap_font_sizes: bool = False
    vectorize_shapes: bool = False
    bitmap_extractor: BitmapExtractor = BitmapExtractor.OPENCV


class JobResponse(BaseModel):
    id: str
    document: DocumentSpec
    original_url: str
    background_url: str


class Capabilities(BaseModel):
    ocr_available: bool
    lama_installed: bool
    devices: list[str]
    fonts: list[str]