from typing import Annotated, Literal

from pydantic import BaseModel, ConfigDict, Field

from olmoearth_shared.models.datasets.collection import Collection


class BaseDataProviderItemProperties(BaseModel):
    """Base class for data provider-specific properties"""

    model_config = ConfigDict(extra="forbid")


class Sentinel1DataProviderItemProperties(BaseDataProviderItemProperties):
    """Data provider properties for Sentinel-1 imagery"""

    collection: Literal[Collection.SENTINEL_1_RTC] = Field(description="Collection identifier")


class Sentinel2DataProviderItemProperties(BaseDataProviderItemProperties):
    """Data provider properties for Sentinel-2 imagery"""

    collection: Literal[Collection.SENTINEL_2_L2A] = Field(description="Collection identifier")
    processing_baseline: str = Field(description="Processing baseline version (e.g., '02.12', '05.10')")
    boa_offset_applied: bool | None = Field(
        default=None,
        description="Whether BOA additive offset corrections were applied (if known)",
    )


class LandsatC2DataProviderItemProperties(BaseDataProviderItemProperties):
    """Data provider properties for Landsat Collection 2 imagery"""

    collection: Literal[Collection.LANDSAT_8_9_C2_L1, Collection.LANDSAT_8_9_C2_L2] = Field(description="Collection identifier")
    processing_correction_level: str = Field(description="Processing/correction level (e.g., 'L2SP')")


DataProviderItemProperties = Annotated[
    Sentinel1DataProviderItemProperties | Sentinel2DataProviderItemProperties | LandsatC2DataProviderItemProperties,
    Field(discriminator="collection"),
]
