from datetime import timedelta
from typing import assert_never

from rslearn.config import dataset

from olmoearth_run.runner.models.operational_context import DataSource, DatasetOperationalContext
from olmoearth_run.runner.tools.olmoearth_config.constants import (
    OLMOEARTH_MODALITIES_SPEC,
    STAC_EO_CLOUD_COVER,
    STAC_SAR_POLARIZATIONS,
    STAC_SAT_ORBIT_STATE_ASCENDING,
    STAC_SAT_ORBIT_STATE_DESCENDING,
)
from olmoearth_run.runner.tools.olmoearth_config.utils import ModalityUtils, TemporalityUtils
from olmoearth_shared.models.olmoearth_config.data.modalities import (
    Modality,
    Sentinel1RTC,
    Sentinel1RTCOrbitDirection,
    Sentinel2L2A,
    Sentinel2L2ASortBy,
    SpaceMode,
    SpaceMode as ModalitiesSpaceMode,
)
from olmoearth_shared.models.olmoearth_config.data.temporality import (
    FramedPointInTimeTemporality,
    ObservationStrategy,
    PeriodTemporality,
    PointInTimeTemporality,
    RepeatingIntervalObservation,
    SampledObservation,
)
from olmoearth_shared.models.olmoearth_config.olmoearth_config import OlmoEarthConfig


class ModalityLayers:
    @staticmethod
    def generate_modality_layers(
        olmoearth_config: OlmoEarthConfig,
        ops_context: DatasetOperationalContext,
    ) -> dict[str, dataset.LayerConfig]:
        modality_layers: dict[str, dataset.LayerConfig] = {}

        if olmoearth_config.data.modalities.sentinel2_l2a:
            modality_layers.update(
                ModalityLayers.generate_layers_for_modality(
                    olmoearth_config.data.modalities.sentinel2_l2a,
                    olmoearth_config,
                    ops_context,
                )
            )

        if olmoearth_config.data.modalities.sentinel1_rtc:
            modality_layers.update(
                ModalityLayers.generate_layers_for_modality(
                    olmoearth_config.data.modalities.sentinel1_rtc,
                    olmoearth_config,
                    ops_context,
                )
            )

        return modality_layers

    @staticmethod
    def generate_layers_for_modality(
        modality: Modality,
        olmoearth_config: OlmoEarthConfig,
        ops_context: DatasetOperationalContext,
    ) -> dict[str, dataset.LayerConfig]:
        """Generate layer configs for a modality based on temporality."""
        temporality = olmoearth_config.data.temporality

        match temporality:
            case PeriodTemporality():
                return ModalityLayers.generate_period_temporality_layers_for_modality(
                    temporality, modality, ops_context
                )
            case PointInTimeTemporality():
                return ModalityLayers.generate_point_in_time_temporality_layers_for_modality(
                    temporality, modality, ops_context
                )
            case FramedPointInTimeTemporality():
                return ModalityLayers.generate_framed_point_in_time_temporality_layers_for_modality(
                    temporality, modality, ops_context
                )
            case _:
                assert_never(temporality)

    @staticmethod
    def generate_period_temporality_layers_for_modality(
        temporality: PeriodTemporality,
        modality: Modality,
        ops_context: DatasetOperationalContext,
    ) -> dict[str, dataset.LayerConfig]:
        """Generate layer config for a modality with period temporality."""
        query_config = DataSourceUtils.generate_query_config_for_observation(
            temporality.observation, modality.space_mode
        )
        duration = TemporalityUtils.generate_duration_for_observation(temporality.observation)
        time_offset = None
        data_source_config = DataSourceUtils.generate_data_source_config_for_modality(
            modality, query_config, ops_context, duration, time_offset
        )
        layer_config = ModalityLayers.generate_layer_config_for_modality(modality, data_source_config)
        return {ModalityUtils.get_layer_name(modality, OLMOEARTH_MODALITIES_SPEC): layer_config}

    @staticmethod
    def generate_point_in_time_temporality_layers_for_modality(
        temporality: PointInTimeTemporality,
        modality: Modality,
        ops_context: DatasetOperationalContext,
    ) -> dict[str, dataset.LayerConfig]:
        """Generate layer config for a modality with point-in-time temporality."""
        query_config = dataset.QueryConfig(
            space_mode=dataset.SpaceMode.CONTAINS,
            min_matches=1,
            max_matches=1,
            period_duration=timedelta(seconds=temporality.observation_window_seconds),
        )
        duration = timedelta(seconds=temporality.observation_window_seconds)
        # Center the observation window on the prediction point
        time_offset = timedelta(seconds=-0.5 * temporality.observation_window_seconds)
        data_source_config = DataSourceUtils.generate_data_source_config_for_modality(
            modality, query_config, ops_context, duration, time_offset
        )
        layer_config = ModalityLayers.generate_layer_config_for_modality(modality, data_source_config)
        return {ModalityUtils.get_layer_name(modality, OLMOEARTH_MODALITIES_SPEC): layer_config}

    @staticmethod
    def generate_framed_point_in_time_temporality_layers_for_modality(
        temporality: FramedPointInTimeTemporality,
        modality: Modality,
        ops_context: DatasetOperationalContext,
    ) -> dict[str, dataset.LayerConfig]:
        """Generate layer config for a modality with framed point-in-time temporality."""
        modality_layers: dict[str, dataset.LayerConfig] = {}

        if temporality.lookbehind_observation:
            lookbehind_query_config = DataSourceUtils.generate_query_config_for_observation(
                temporality.lookbehind_observation, modality.space_mode
            )
            lookbehind_duration = TemporalityUtils.generate_duration_for_observation(
                temporality.lookbehind_observation
            )
            # Calculate offset so window ends at T - lookbehind_offset_days
            lookbehind_offset_days = temporality.lookbehind_offset_days or 0
            lookbehind_time_offset = -timedelta(days=lookbehind_offset_days) - lookbehind_duration
            lookbehind_data_source_config = DataSourceUtils.generate_data_source_config_for_modality(
                modality, lookbehind_query_config, ops_context, lookbehind_duration, lookbehind_time_offset
            )
            lookbehind_layer_config = ModalityLayers.generate_layer_config_for_modality(
                modality, lookbehind_data_source_config
            )
            modality_layers[ModalityUtils.get_lookbehind_layer_name(modality, OLMOEARTH_MODALITIES_SPEC)] = lookbehind_layer_config

        if temporality.lookahead_observation:
            lookahead_query_config = DataSourceUtils.generate_query_config_for_observation(
                temporality.lookahead_observation, modality.space_mode
            )
            lookahead_duration = TemporalityUtils.generate_duration_for_observation(
                temporality.lookahead_observation
            )
            # Calculate offset so window starts at T + lookahead_offset_days
            lookahead_time_offset = (
                timedelta(days=temporality.lookahead_offset_days)
                if temporality.lookahead_offset_days is not None
                else None
            )
            lookahead_data_source_config = DataSourceUtils.generate_data_source_config_for_modality(
                modality, lookahead_query_config, ops_context, lookahead_duration, lookahead_time_offset
            )
            lookahead_layer_config = ModalityLayers.generate_layer_config_for_modality(
                modality, lookahead_data_source_config
            )
            modality_layers[ModalityUtils.get_lookahead_layer_name(modality, OLMOEARTH_MODALITIES_SPEC)] = lookahead_layer_config

        return modality_layers

    @staticmethod
    def generate_layer_config_for_modality(
        modality: Modality,
        data_source_config: dataset.DataSourceConfig,
    ) -> dataset.LayerConfig:
        """Generate a LayerConfig for a given modality."""
        match modality:
            case Sentinel2L2A():
                return ModalityLayers.generate_layer_config_for_s2_l2a(data_source_config)
            case Sentinel1RTC():
                return ModalityLayers.generate_layer_config_for_sentinel1_rtc(data_source_config)
            case _:
                assert_never(modality)

    @staticmethod
    def generate_layer_config_for_sentinel1_rtc(
        data_source_config: dataset.DataSourceConfig,
    ) -> dataset.LayerConfig:
        return dataset.LayerConfig(
            type=dataset.LayerType.RASTER,
            data_source=data_source_config,
            band_sets=[
                dataset.BandSetConfig(
                    bands=OLMOEARTH_MODALITIES_SPEC.sentinel1_rtc_band_order,
                    dtype=dataset.DType.FLOAT32,
                )
            ],
        )

    @staticmethod
    def generate_layer_config_for_s2_l2a(
        data_source_config: dataset.DataSourceConfig,
    ) -> dataset.LayerConfig:
        return dataset.LayerConfig(
            type=dataset.LayerType.RASTER,
            data_source=data_source_config,
            band_sets=[
                dataset.BandSetConfig(
                    bands=OLMOEARTH_MODALITIES_SPEC.sentinel2_l2a_band_order,
                    dtype=dataset.DType.UINT16,
                )
            ],
        )


class DataSourceUtils:
    @staticmethod
    def generate_query_config_for_observation(
        observation_strategy: ObservationStrategy,
        space_mode: SpaceMode,
    ) -> dataset.QueryConfig:
        """Generate a QueryConfig for a given observation strategy."""
        num_matches = TemporalityUtils.find_num_matches_for_observation(observation_strategy)
        return dataset.QueryConfig(
            space_mode=ConfigMappers.map_space_mode_from_observation_strategy(
                space_mode, observation_strategy
            ),
            min_matches=num_matches,
            max_matches=num_matches,
            period_duration=TemporalityUtils.generate_period_duration_for_observation(
                observation_strategy
            ),
        )

    @staticmethod
    def generate_data_source_config_for_modality(
        modality: Modality,
        query_config: dataset.QueryConfig,
        ops_context: DatasetOperationalContext,
        duration: timedelta,
        time_offset: timedelta | None,
    ) -> dataset.DataSourceConfig:
        """Generate a DataSourceConfig for a given modality and data source."""
        match ops_context.data_source:
            case DataSource.PLANETARY_COMPUTER:
                match modality:
                    case Sentinel2L2A():
                        return DataSourceUtils.generate_planetary_computer_data_source_config_for_s2_l2a(
                            modality, query_config, ops_context, duration, time_offset
                        )
                    case Sentinel1RTC():
                        return DataSourceUtils.generate_planetary_computer_data_source_config_for_sentinel1_rtc(
                            modality, query_config, ops_context, duration, time_offset
                        )
                    case _:
                        assert_never(modality)
            case _:
                assert_never(ops_context.data_source)

    @staticmethod
    def generate_planetary_computer_data_source_config_for_sentinel1_rtc(
        rtc_config: Sentinel1RTC,
        query_config: dataset.QueryConfig,
        ops_context: DatasetOperationalContext,
        duration: timedelta,
        time_offset: timedelta | None,
    ) -> dataset.DataSourceConfig:
        return dataset.DataSourceConfig(
            class_path="rslearn.data_sources.planetary_computer.Sentinel1",
            init_args=dict(
                query={
                    "sar:instrument_mode": {"eq": "IW"},
                    "sar:polarizations": {"eq": STAC_SAR_POLARIZATIONS},
                    "sat:orbit_state": {
                        "eq": ConfigMappers.map_sentinel1_rtc_orbit_direction(rtc_config.orbit_direction)
                    },
                },
                cache_dir="cache/planetary_computer",
                timeout="0:0:10",  # in h:m:s format required by rslearn's jsonargparse
            ),
            query_config=query_config,
            duration=duration,
            time_offset=time_offset,
            ingest=False,
        )

    @staticmethod
    def generate_planetary_computer_data_source_config_for_s2_l2a(
        s2_config: Sentinel2L2A,
        query_config: dataset.QueryConfig,
        ops_context: DatasetOperationalContext,
        duration: timedelta,
        time_offset: timedelta | None,
    ) -> dataset.DataSourceConfig:
        return dataset.DataSourceConfig(
            class_path="rslearn.data_sources.planetary_computer.Sentinel2",
            init_args=dict(
                harmonize=True,
                sort_by=ConfigMappers.map_s2_sort_by(s2_config.sort_by),
                cache_dir="cache/planetary_computer",
                timeout="0:0:10",  # in h:m:s format required by rslearn's jsonargparse
            ),
            query_config=query_config,
            duration=duration,
            time_offset=time_offset,
            ingest=False,
        )


class ConfigMappers:
    """Maps simplified unified config values into rslearn settings."""

    @staticmethod
    def map_space_mode_from_observation_strategy(
        space_mode: ModalitiesSpaceMode,
        observation_strategy: ObservationStrategy,
    ) -> dataset.SpaceMode:
        match space_mode:
            case ModalitiesSpaceMode.CONTAINS:
                return dataset.SpaceMode.CONTAINS
            case ModalitiesSpaceMode.MOSAIC:
                match observation_strategy:
                    case SampledObservation():
                        return dataset.SpaceMode.MOSAIC
                    case RepeatingIntervalObservation():
                        return dataset.SpaceMode.PER_PERIOD_MOSAIC
                    case _:
                        assert_never(observation_strategy)
            case _:
                assert_never(space_mode)

    @staticmethod
    def map_s2_sort_by(sort_by: Sentinel2L2ASortBy) -> str:
        match sort_by:
            case Sentinel2L2ASortBy.CLOUD_COVER:
                return STAC_EO_CLOUD_COVER
            case _:
                assert_never(sort_by)

    @staticmethod
    def map_sentinel1_rtc_orbit_direction(orbit_direction: Sentinel1RTCOrbitDirection) -> str:
        match orbit_direction:
            case Sentinel1RTCOrbitDirection.ASCENDING_ORBIT:
                return STAC_SAT_ORBIT_STATE_ASCENDING
            case Sentinel1RTCOrbitDirection.DESCENDING_ORBIT:
                return STAC_SAT_ORBIT_STATE_DESCENDING
            case _:
                assert_never(orbit_direction)
