from enum import StrEnum
from typing import Annotated, Literal

from pydantic import BaseModel, Field
from typing_extensions import TypeAliasType


class WindowPreparerName(StrEnum):
    POINT_TO_RASTER = "point_to_raster"
    POLYGON_TO_RASTER = "polygon_to_raster"


class GriddingStrategy(StrEnum):
    """Strategy for dividing annotation tasks into grid cells.

    CENTER_CROP: Crop to the largest centered rectangle that fits whole cells.

    Example: 120×120 pixel raster with 50×50 pixel cells

    120 ÷ 50 = 2 whole cells + 20px remainder
    Remainder split equally: 10px margin on each side

            10px    50px      50px    10px
            ┌─────┬──────────┬──────────┬─────┐
       10px │░░░░░│░░░░░░░░░░│░░░░░░░░░░│░░░░░│
            ├─────┼──────────┼──────────┼─────┤
            │░░░░░│          │          │░░░░░│
       50px │░░░░░│  Cell 1  │  Cell 2  │░░░░░│
            │░░░░░│          │          │░░░░░│
            ├─────┼──────────┼──────────┼─────┤
            │░░░░░│          │          │░░░░░│
       50px │░░░░░│  Cell 3  │  Cell 4  │░░░░░│
            │░░░░░│          │          │░░░░░│
            ├─────┼──────────┼──────────┼─────┤
       10px │░░░░░│░░░░░░░░░░│░░░░░░░░░░│░░░░░│
            └─────┴──────────┴──────────┴─────┘

            ░░░░░ = NODATA (not included in any cell)

    Result: 4 cells (2×2), each with complete valid data
            Outer 10px border becomes nodata

    ==============================

    PAD_TO_FILL: Extend edge cells beyond task bounds (nodata fills naturally).

    Example: 120×120 pixel raster with 50×50 pixel cells

    ceil(120 ÷ 50) = 3 cells needed per dimension
    Grid extends 30px beyond raster on right and bottom

         ←─ 50px ─→←─ 50px ─→←─ 50px ─→
                             ←20→←─30─→
                            valid nodata
        ┌──────────┬──────────┬────┬─────┐  ↑
        │          │          │    │░░░░░│  │
        │  Cell 1  │  Cell 2  │    │░░3░░│ 50px
        │          │          │    │░░░░░│  │
        ├──────────┼──────────┼────┼─────┤  ↓
        │          │          │    │░░░░░│  ↑
        │  Cell 4  │  Cell 5  │    │░░6░░│ 50px
        │          │          │    │░░░░░│  │
        ├──────────┼──────────┼────┼─────┤  ↓  ↑ 20px valid
        │  Cell 7  │  Cell 8  │    │░░9░░│     ↓
        ├──────────┼──────────┼────┼─────┤     ↑ 30px nodata
        │░░░░░░░░░░│░░░░░░░░░░│░░░░│░░░░░│     ↓
        └──────────┴──────────┴────┴─────┘

        ░░░░░ = NODATA

    Cells 1, 2, 4, 5:  100% valid data
    Cells 3, 6:        rightmost 30px = nodata
    Cells 7, 8:        bottom 30px = nodata
    Cell 9:            right 30px AND bottom 30px = nodata
    """

    CENTER_CROP = "center_crop"
    PAD_TO_FILL = "pad_to_fill"


class PointToRasterWindowPreparer(BaseModel):
    """Preparer that converts point annotations into raster windows centered on the point."""

    name: Literal[WindowPreparerName.POINT_TO_RASTER]
    window_buffer: int = Field(ge=0, description="The buffer around the point in pixels.")


class PolygonToRasterWindowPreparer(BaseModel):
    """Preparer that converts polygon annotations into raster windows.

    By default, creates one window per annotation task using the task geometry bounds.
    When gridding_strategy is set, divides tasks into fixed-size grid cells (size from input_preprocessing.input_size).
    """

    name: Literal[WindowPreparerName.POLYGON_TO_RASTER]
    gridding_strategy: GriddingStrategy | None = Field(
        default=None,
        description="If set, enables grid mode and specifies how to handle non-integer divisions.",
    )
    drop_inconsistent_utm: bool = Field(
        default=False,
        description="If True, skip tasks where corners span different UTM zones. By default chooses utm based on task centroid."
    )


WindowPreparer = TypeAliasType(
    "WindowPreparer",
    Annotated[
        PointToRasterWindowPreparer | PolygonToRasterWindowPreparer,
        Field(discriminator="name")
    ]
)
