from olmoearth_shared.models.datasets.bands.landsat import LandsatC2Band
from olmoearth_shared.models.datasets.bands.sentinel1 import Sentinel1Band
from olmoearth_shared.models.datasets.bands.sentinel2 import Sentinel2L2ABand
from olmoearth_shared.models.datasets.collection import Collection

BandEnum = type[Sentinel1Band] | type[Sentinel2L2ABand] | type[LandsatC2Band]

COLLECTION_BANDS: dict[Collection, BandEnum] = {
    Collection.SENTINEL_1_RTC: Sentinel1Band,
    Collection.SENTINEL_2_L2A: Sentinel2L2ABand,
    Collection.LANDSAT_8_9_C2_L1: LandsatC2Band,
    Collection.LANDSAT_8_9_C2_L2: LandsatC2Band,
}


def get_band_enum(collection: Collection) -> BandEnum:
    """Get the band enum class for a collection.

    Args:
        collection: The satellite imagery collection.

    Returns:
        The band enum class for the collection.

    Raises:
        KeyError: If the collection is not supported.
    """
    return COLLECTION_BANDS[collection]


def get_valid_bands(collection: Collection) -> set[str]:
    """Get the set of valid band names for a collection.

    Args:
        collection: The satellite imagery collection.

    Returns:
        Set of valid band name strings for the collection.
    """
    band_enum = get_band_enum(collection)
    return {band.value for band in band_enum}


def validate_band(band: str, collection: Collection) -> None:
    """Validate that a band name is valid for a collection.

    Args:
        band: The band name to validate.
        collection: The satellite imagery collection.

    Raises:
        ValueError: If the band is not valid for the collection.
    """
    valid_bands = get_valid_bands(collection)
    if band not in valid_bands:
        valid_str = ", ".join(sorted(valid_bands))
        raise ValueError(f"Invalid band '{band}' for {collection.value}. Valid bands: {valid_str}")
