from datetime import datetime
from typing import Annotated, Literal

from geojson_pydantic.geometries import Geometry
from pydantic import BaseModel, ConfigDict, Field

from olmoearth_shared.models.datasets.collection import Collection
from olmoearth_shared.models.datasets.instrument_mode import InstrumentMode
from olmoearth_shared.models.datasets.landsat_collection_category import LandsatCollectionCategory
from olmoearth_shared.models.datasets.orbit_direction import OrbitDirection
from olmoearth_shared.models.datasets.polarization import Polarization
from olmoearth_shared.models.datasets.satellite import Satellite


class BaseItemProperties(BaseModel):
    """Base properties shared by all items"""

    model_config = ConfigDict(extra="forbid")

    created_at: datetime = Field(description="When this item was created in the system")
    updated_at: datetime = Field(description="When this item was last updated in the system")
    collected_at: datetime = Field(description="Collection time")
    geometry: Geometry = Field(description="The geometry of the scene")
    satellite: Satellite = Field(description="Satellite identifier")


class Sentinel1ItemProperties(BaseItemProperties):
    """Properties for Sentinel-1 imagery"""

    collection: Literal[Collection.SENTINEL_1_RTC] = Field(description="Collection identifier")
    orbit_direction: OrbitDirection = Field(description="Orbit direction")
    relative_orbit: int = Field(description="Relative orbit number")
    instrument_mode: InstrumentMode = Field(description="SAR instrument mode")
    polarizations: set[Polarization] = Field(description="Polarization configuration")


class Sentinel2ItemProperties(BaseItemProperties):
    """Properties for Sentinel-2 imagery"""

    collection: Literal[Collection.SENTINEL_2_L2A] = Field(description="Collection identifier")
    cloud_cover: float = Field(description="Cloud cover percentage")
    relative_orbit: int = Field(description="Relative orbit number")
    mgrs_tile: str = Field(description="MGRS tile identifier")


class LandsatC2ItemProperties(BaseItemProperties):
    """Properties for Landsat Collection 2 imagery"""

    collection: Literal[Collection.LANDSAT_8_9_C2_L1, Collection.LANDSAT_8_9_C2_L2] = Field(description="Collection identifier")
    cloud_cover: float | None = Field(description="Cloud cover percentage")
    wrs_path: int = Field(description="WRS-2 path")
    wrs_row: int = Field(description="WRS-2 row")
    collection_category: LandsatCollectionCategory = Field(description="Collection category (RT, T1, or T2)")


ItemProperties = Annotated[
    Sentinel1ItemProperties | Sentinel2ItemProperties | LandsatC2ItemProperties,
    Field(discriminator="collection"),
]
