"""rslearn data source for olmoearth_datasets.

This makes search requests on olmoearth_datasets to get matching items. The items
specify which data providers (Planetary Computer, miscellaneous AWS S3 buckets, etc.)
assets can be downloaded from. During ingestion or materialization, we try to download
image data from each provider, falling back to the next provider when we encounter
errors.
"""

import json
import os
import random
import tempfile
from collections.abc import Callable, Iterator
from datetime import timedelta
from typing import Any
from typing_extensions import override

from geojson_pydantic.geometries import parse_geometry_obj
import numpy.typing as npt
import rasterio  # type: ignore[import-untyped]
import shapely
from rasterio.enums import Resampling  # type: ignore[import-untyped]
from upath import UPath

from olmoearth_shared.api.common.search_filters import (
    DatetimeFilter,
    KeywordFilter,
)
from olmoearth_shared.http.request_session import RequestSession
from olmoearth_shared.api.datasets.items import SearchItemsRequest
from olmoearth_shared.models.datasets.asset import Asset
from olmoearth_shared.models.datasets.collection import Collection
from olmoearth_shared.models.datasets.data_provider_item import DataProviderItem
from olmoearth_shared.models.datasets.data_provider_name import DataProviderName
from olmoearth_shared.models.datasets.item import Item as ApiItem
from olmoearth_shared.models.datasets.bands.validation import validate_band
from rslearn.config import QueryConfig
from rslearn.const import WGS84_PROJECTION
from rslearn.data_sources.data_source import DataSourceContext, Item
from rslearn.data_sources.direct_materialize_data_source import (
    DirectMaterializeDataSource,
)
from rslearn.data_sources.utils import match_candidate_items_to_window
from rslearn.log_utils import get_logger
from rslearn.tile_stores import TileStoreWithLayer
from rslearn.utils.geometry import PixelBounds, Projection, STGeometry
from rslearn.utils.raster_format import get_raster_projection_and_bounds

from olmoearth_run.config import OlmoEarthSettings
from olmoearth_run.runner.metrics.olmoearth_datasets import (
    OlmoEarthDatasetsMetrics,
    Outcome,
)

from .retry_config import RetryConfig
from .search_items_query_filters import SearchItemsQueryFilters

logger = get_logger(__name__)


def select_band_sets(
    available_band_sets: list[list[str]],
    band_sets: list[list[str]] | None,
    context: DataSourceContext,
) -> list[list[str]]:
    """Select band sets based on context/explicit argument.

    This helper determines which band sets to use for a data source based on:
    1. If context has a layer_config, select band sets that intersect with configured bands
    2. If band_sets argument is provided, use those directly
    3. Otherwise, use all available band sets

    Args:
        available_band_sets: all band sets available for the modality
        band_sets: explicit band sets argument from caller
        context: the data source context

    Returns:
        The selected list of band sets to use.
    """
    if context.layer_config is not None:
        # Select band sets that intersect with configured bands.
        selected: list[list[str]] = []
        for bands in available_band_sets:
            for band_set in context.layer_config.band_sets:
                if not set(band_set.bands).intersection(set(bands)):
                    continue
                selected.append(bands)
                break
        return selected
    elif band_sets is not None:
        return band_sets
    else:
        return available_band_sets


class OlmoEarthDatasets(DirectMaterializeDataSource[Item]):
    """rslearn data source for the olmoearth_datasets API.

    This data source queries the olmoearth_datasets API for scene metadata, and
    downloads assets from the various data providers that it returns.

    It can be configured directly, but most users will use convenience subclasses
    which are provided for each collection.

    When downloading an asset, if an item has multiple data providers, the providers
    are shuffled and tried in sequence until one succeeds. This provides automatic
    fallback if one provider is unavailable.
    """

    def __init__(
        self,
        collection: Collection,
        band_sets: list[list[str]],
        query: SearchItemsQueryFilters | None = None,
        retry_config: RetryConfig | None = None,
        timeout: timedelta = timedelta(seconds=30),
        skip_items_missing_assets: bool = True,
        context: DataSourceContext = DataSourceContext(),
    ):
        """Initialize a new OlmoEarthDatasets instance.

        Args:
            collection: the collection to query (e.g., Collection.SENTINEL_2_L2A).
            band_sets: list of band sets to ingest/materialize. Each band set must be a
                list of band names in a single asset. Note that we assume that the
                groups of band names are consistent across data providers, although
                each provider may use different asset names.
            query: query filters and sorting options to include in API requests.
            retry_config: configuration for HTTP request retries.
            timeout: timeout for API and download requests.
            skip_items_missing_assets: skip items that are missing any of the assets
                in band_sets during get_items. Useful for filtering out items that
                don't have all required bands (e.g., single-polarization Sentinel-1).
            context: the data source context.
        """
        # Validate provided bands.
        for bands in band_sets:
            for band in bands:
                validate_band(band, collection)

        # Construct asset_bands for parent class (DirectMaterializeDataSource): we need to give
        # it asset keys, but the actual asset names vary across providers, so here we
        # compute suitable asset keys based on the band names.
        asset_bands = {",".join(bands): bands for bands in band_sets}
        DirectMaterializeDataSource.__init__(self, asset_bands=asset_bands)

        self.band_sets = band_sets
        self.collection = collection
        self.query = query if query is not None else SearchItemsQueryFilters()
        self.timeout = timeout
        self.skip_items_missing_assets = skip_items_missing_assets
        self.context = context

        if retry_config is None:
            retry_config = RetryConfig()
        self._session = RequestSession(
            total_retries=retry_config.total_retries,
            backoff_factor=retry_config.backoff_factor,
            backoff_max=retry_config.backoff_max,
        )

    def _search_request(
        self,
        item_id: str | None = None,
        geometry: STGeometry | None = None,
    ) -> Iterator[ApiItem]:
        """Make a search request to olmoearth_datasets.

        Automatically paginates through all results, yielding each ApiItem as it's
        fetched.

        Args:
            item_id: filter by exact item ID.
            geometry: STGeometry to filter by spatial and temporal intersection.

        Yields:
            ApiItem objects from the API.
        """
        # Build the search request using SearchItemsRequest.
        page_size = 1000
        request = SearchItemsRequest(
            collection=KeywordFilter[Collection](eq=self.collection),
            id=KeywordFilter[str](eq=item_id) if item_id else None,
            # Pagination arguments -- we will update these in the pagination loop.
            limit=page_size,
            offset=0,
            # Ask the API sign URLs for us, so we can handle URLs from different
            # providers in the same way here.
            sign_urls=True,
            # Apply query filters.
            sort_by=self.query.sort_by,
            sort_direction=self.query.sort_direction,
            satellite=self.query.satellite,
            cloud_cover=self.query.cloud_cover,
            mgrs_tile=self.query.mgrs_tile,
            wrs_path=self.query.wrs_path,
            wrs_row=self.query.wrs_row,
            collection_category=self.query.collection_category,
            orbit_direction=self.query.orbit_direction,
            relative_orbit=self.query.relative_orbit,
            instrument_mode=self.query.instrument_mode,
            polarization=self.query.polarizations,
        )

        # Extract spatial extent and time range from STGeometry.
        if geometry is not None:
            wgs84_geometry = geometry.to_projection(WGS84_PROJECTION)
            geojson_dict = json.loads(shapely.to_geojson(wgs84_geometry.shp))
            request.intersects_geometry = parse_geometry_obj(geojson_dict)

            if wgs84_geometry.time_range is not None:
                request.collected_at = DatetimeFilter(
                    gte=wgs84_geometry.time_range[0],
                    lt=wgs84_geometry.time_range[1],
                )

        # Paginate through all results.
        api_base_url = OlmoEarthSettings.OEDATASETS_API_URL.rstrip("/")
        url = f"{api_base_url}/api/v1/items/search"

        while True:
            response = self._session.post(
                url,
                json=request.model_dump(mode="json", exclude_none=True),
                headers={"Content-Type": "application/json"},
                timeout=self.timeout.total_seconds(),
            )
            response.raise_for_status()
            records = response.json()["records"]

            logger.debug(
                "OlmoEarthDatasets: got %d records at offset %d",
                len(records),
                request.offset,
            )
            for record in records:
                yield ApiItem.model_validate(record)

            # Stop if we got fewer than requested (no more results).
            if len(records) < page_size:
                break

            request.offset += len(records)

    def _api_item_to_rslearn_item(self, api_item: ApiItem) -> Item:
        """Convert an ApiItem to an rslearn Item."""
        # Extract geometry for rslearn's STGeometry format.
        # We use collected_at as both start and end of time range.
        shp = shapely.geometry.shape(api_item.properties.geometry)
        collected_at = api_item.properties.collected_at
        time_range = (collected_at, collected_at)
        geometry = STGeometry(WGS84_PROJECTION, shp, time_range)

        return Item(name=api_item.id, geometry=geometry)

    def _item_has_all_bands(self, api_item: ApiItem) -> bool:
        """Check if an item has all required band sets from at least one provider."""
        for bands in self.band_sets:
            if not self._get_providers_with_asset(api_item, bands):
                return False
        return True

    @override
    def get_items(
        self, geometries: list[STGeometry], query_config: QueryConfig
    ) -> list[list[list[Item]]]:
        """Get a list of items in the data source intersecting the given geometries.

        Args:
            geometries: the spatiotemporal geometries.
            query_config: the query configuration.

        Returns:
            List of groups of items that should be retrieved for each geometry.
        """
        groups = []
        for geometry in geometries:
            api_items = list(self._search_request(geometry=geometry))

            # Filter out items missing required assets if configured.
            if self.skip_items_missing_assets:
                original_count = len(api_items)
                api_items = [
                    api_item
                    for api_item in api_items
                    if self._item_has_all_bands(api_item)
                ]
                logger.debug(
                    "skip_items_missing_assets filtered from %d to %d items",
                    original_count,
                    len(api_items),
                )

            # Convert to rslearn Items.
            candidate_items = [
                self._api_item_to_rslearn_item(api_item) for api_item in api_items
            ]

            # Use the standard matching logic.
            cur_groups = match_candidate_items_to_window(
                geometry, candidate_items, query_config
            )
            groups.append(cur_groups)

        return groups

    def _get_api_item_by_name(self, name: str) -> ApiItem:
        """Gets an ApiItem by name from the API."""
        api_items = list(self._search_request(item_id=name))
        if len(api_items) != 1:
            raise ValueError(
                f"Expected 1 item for {name}, got {len(api_items)} from OlmoEarth API"
            )
        return api_items[0]

    @override
    def get_item_by_name(self, name: str) -> Item:
        """Gets an item by name.

        Args:
            name: the item ID. Examples:
                - Sentinel-2 L2A: "S2B_MSIL2A_20250115T185649_R113_T10SEH"
                - Sentinel-1 RTC: "S1A_IW_GRDH_1SDV_20250115T141459_20250115T141524_057539_071528"
                - Landsat C2L1: "LC09_L1TP_046028_20250112_02_T1"

        Returns:
            the Item.
        """
        api_item = self._get_api_item_by_name(name)
        return self._api_item_to_rslearn_item(api_item)

    @override
    def deserialize_item(self, serialized_item: dict) -> Item:
        """Deserializes an item from JSON-decoded data."""
        return Item.deserialize(serialized_item)

    def _download_asset(
        self,
        asset: Asset,
        dest_path: str,
    ) -> None:
        """Download an asset from a provider.

        Args:
            asset: the asset to download.
            dest_path: the local file path to save to.
        """
        logger.debug(
            "Downloading asset %s from %s",
            asset.name,
            asset.url,
        )
        with self._session.get(
            asset.url, stream=True, timeout=self.timeout.total_seconds()
        ) as r:
            r.raise_for_status()
            with open(dest_path, "wb") as f:
                for chunk in r.iter_content(chunk_size=8192):
                    f.write(chunk)

    def _get_providers_with_asset(
        self,
        api_item: ApiItem,
        bands: list[str],
    ) -> list[tuple[DataProviderName, DataProviderItem, Asset]]:
        """Get all providers that have an asset with the specified bands.

        Args:
            api_item: the ApiItem from the API.
            bands: the band names to look for. Matches assets by bands rather than
                asset name since different providers use different names (e.g.,
                "B04" vs "red") but report the same bands.

        Returns:
            list of (provider_name, provider, asset) tuples that have the asset.
        """
        results: list[tuple[DataProviderName, DataProviderItem, Asset]] = []
        for provider_name, provider in api_item.data_providers.items():
            for asset in provider.assets:
                if set(asset.bands) == set(bands):
                    results.append((provider_name, provider, asset))
                    break
        return results

    def _ingest_raster(
        self,
        tile_store: TileStoreWithLayer,
        item_name: str,
        provider_name: str,
        provider: DataProviderItem,
        asset: Asset,
    ) -> None:
        """Download and ingest a single raster asset.

        Args:
            tile_store: the tile store to ingest into.
            item_name: the name of the item being ingested.
            provider_name: the name of the data provider to download from.
            provider: the DataProviderItem corresponding to provider_name.
            asset: the asset to download.
        """
        with tempfile.TemporaryDirectory() as tmp_dir:
            local_fname = os.path.join(tmp_dir, f"{asset.name}.tif")
            self._download_asset(asset, local_fname)

            # Apply any post-processing callback.
            callback = self.get_raster_callback(
                item_name, asset.bands, provider_name, provider
            )
            if callback is not None:
                with rasterio.open(local_fname) as src:
                    array = src.read()
                    projection, raster_bounds = get_raster_projection_and_bounds(src)
                array = callback(array)
                tile_store.write_raster(
                    item_name,
                    asset.bands,
                    projection,
                    raster_bounds,
                    array,
                )
            else:
                tile_store.write_raster_file(item_name, asset.bands, UPath(local_fname))

    @override
    def ingest(
        self,
        tile_store: TileStoreWithLayer,
        items: list[Item],
        geometries: list[list[STGeometry]],
    ) -> None:
        """Ingest items into the given tile store.

        For each item and asset, tries providers in shuffled order until one succeeds.
        Does not retry on failure (retries are handled outside the data source).

        Args:
            tile_store: the tile store to ingest into.
            items: the items to ingest.
            geometries: a list of geometries needed for each item.
        """
        for item in items:
            # Get the ApiItem for this Item. We could cache the ApiItem in the Item
            # object, but the asset URLs expire so instead we only store the name and
            # STGeometry in the Item, and get a fresh ApiItem as needed.
            api_item = self._get_api_item_by_name(item.name)

            for bands in self.band_sets:
                if tile_store.is_raster_ready(item.name, bands):
                    continue

                # Get all providers that have this asset.
                providers_with_asset = self._get_providers_with_asset(api_item, bands)
                if not providers_with_asset:
                    logger.warning(
                        "Item %s has no providers with bands %s",
                        item.name,
                        bands,
                    )
                    continue

                # Shuffle providers to distribute load and provide automatic fallback.
                random.shuffle(providers_with_asset)

                # Try each provider in sequence.
                for idx, (provider_name, provider, asset) in enumerate(
                    providers_with_asset
                ):
                    is_last = idx == len(providers_with_asset) - 1
                    try:
                        self._ingest_raster(
                            tile_store, item.name, provider_name, provider, asset
                        )
                        logger.debug(
                            "Successfully ingested item %s asset %s from %s",
                            item.name,
                            asset.name,
                            provider_name,
                        )
                        OlmoEarthDatasetsMetrics.record_ingest_attempt(
                            provider_name, Outcome.SUCCESS
                        )
                        break
                    except Exception:
                        OlmoEarthDatasetsMetrics.record_ingest_attempt(
                            provider_name, Outcome.FAILURE
                        )
                        if is_last:
                            # If we have gone through all the candidate providers, then raise
                            # the error. rslearn logic outside the data source will handle
                            # retries.
                            raise
                        logger.exception(
                            "Failed to download item %s asset %s from %s",
                            item.name,
                            asset.name,
                            provider_name,
                        )

    # --- Hooks for subclasses ---

    def get_raster_callback(
        self,
        item_name: str,
        bands: list[str],
        provider_name: str,
        provider: DataProviderItem,
    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
        """Return a callback to post-process raster data after reading.

        Subclasses can override this to apply transformations to the raw raster data,
        such as harmonization for Sentinel-2 data.

        Args:
            item_name: the name of the item being read.
            bands: the bands being read.
            provider_name: the name of the data provider being used.
            provider: the corresponding DataProviderItem.

        Returns:
            A callback function that takes an array and returns a modified array,
            or None if no post-processing is needed.
        """
        return None

    # --- DirectMaterializeDataSource implementation ---

    @override
    def get_asset_url(self, item_name: str, asset_key: str) -> str:
        """Not implemented since we override read_raster to iterate over providers."""
        raise NotImplementedError(
            "get_asset_url is not implemented, read_raster should be called instead"
        )

    @override
    def read_raster(
        self,
        layer_name: str,
        item_name: str,
        bands: list[str],
        projection: Projection,
        bounds: PixelBounds,
        resampling: Resampling = Resampling.bilinear,
    ) -> npt.NDArray[Any]:
        """Read raster data, trying multiple providers.

        Args:
            layer_name: the layer name or alias.
            item_name: the item to read.
            bands: the list of bands identifying which specific raster to read.
            projection: the projection to read in.
            bounds: the bounds to read.
            resampling: the resampling method to use in case reprojection is needed.

        Returns:
            the raster data as a numpy array.
        """
        api_item = self._get_api_item_by_name(item_name)
        providers_with_asset = self._get_providers_with_asset(api_item, bands)

        if not providers_with_asset:
            raise ValueError(f"No provider has bands {bands} for item {item_name}")

        # Shuffle providers to distribute load and provide automatic fallback.
        random.shuffle(providers_with_asset)

        # Try each provider in sequence.
        for idx, (provider_name, provider, asset) in enumerate(providers_with_asset):
            is_last = idx == len(providers_with_asset) - 1
            try:
                raw_data = self._read_raster_from_url(
                    asset.url, projection, bounds, resampling
                )

                # Apply any post-processing callback.
                callback = self.get_raster_callback(
                    item_name, bands, provider_name, provider
                )
                if callback is not None:
                    raw_data = callback(raw_data)

                return raw_data
            except Exception:
                if is_last:
                    # If we have gone through all the candidate providers, then raise
                    # the error. rslearn logic outside the data source will handle
                    # retries.
                    raise
                logger.exception(
                    "Failed to read item %s asset %s from %s",
                    item_name,
                    asset.name,
                    provider_name,
                )

        # Should never reach here, but satisfy type checker.
        raise RuntimeError("No providers available")

    @override
    def get_raster_bands(self, layer_name: str, item_name: str) -> list[list[str]]:
        """Get the sets of bands that have been stored for the specified item.

        Returns the bands from asset_bands for assets that exist on the item.
        """
        api_item = self._get_api_item_by_name(item_name)
        all_bands = []
        for bands in self.band_sets:
            # Check if any provider has this asset.
            if self._get_providers_with_asset(api_item, bands):
                all_bands.append(bands)
        return all_bands
