"""Sentinel-2 L2A data from olmoearth_datasets."""

from xml.etree.ElementTree import Element

from collections.abc import Callable
from datetime import timedelta
from typing import Any
from typing_extensions import override

import defusedxml.ElementTree as ET
import numpy as np
import numpy.typing as npt
from upath import UPath

from olmoearth_shared.models.datasets.bands.sentinel2 import Sentinel2L2ABand
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_item_properties import (
    Sentinel2DataProviderItemProperties,
)
from rslearn.data_sources.copernicus import get_harmonize_callback
from rslearn.data_sources.data_source import DataSourceContext
from rslearn.utils.fsspec import join_upath


from .olmoearth_datasets import OlmoEarthDatasets, select_band_sets
from .retry_config import RetryConfig
from .search_items_query_filters import SearchItemsQueryFilters


class Sentinel2L2A(OlmoEarthDatasets):
    """Data source for Sentinel-2 L2A data from olmoearth_datasets."""

    COLLECTION = Collection.SENTINEL_2_L2A

    # Available band sets for Sentinel-2. Most bands are single-band assets, but
    # the visual composite (R, G, B) is a multi-band asset (this information is not in
    # Sentinel2L2ABand so we have to enumerate the assets here).
    # B01-B8A require harmonization while SCL and RGB do not.
    BAND_SETS_NEEDING_HARMONIZATION: list[list[str]] = [
        [Sentinel2L2ABand.B01],
        [Sentinel2L2ABand.B02],
        [Sentinel2L2ABand.B03],
        [Sentinel2L2ABand.B04],
        [Sentinel2L2ABand.B05],
        [Sentinel2L2ABand.B06],
        [Sentinel2L2ABand.B07],
        [Sentinel2L2ABand.B08],
        [Sentinel2L2ABand.B09],
        [Sentinel2L2ABand.B11],
        [Sentinel2L2ABand.B12],
        [Sentinel2L2ABand.B8A],
    ]
    BAND_SETS_NOT_NEEDING_HARMONIZATION: list[list[str]] = [
        [Sentinel2L2ABand.R, Sentinel2L2ABand.G, Sentinel2L2ABand.B],
        [Sentinel2L2ABand.SCL],
    ]

    # Standard BOA offset applied in processing baseline 04.00+.
    HARMONIZE_OFFSET = -1000

    # Asset name for product metadata XML (used for harmonization).
    PRODUCT_METADATA_ASSET = "product-metadata"

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

        Args:
            harmonize: harmonize pixel values across different processing baselines.
                See https://developers.google.com/earth-engine/datasets/catalog/COPERNICUS_S2_SR_HARMONIZED
            band_sets: list of band sets to ingest. If None, band sets are derived from
                the layer config (if available), otherwise all default band sets.
            query: query filters and sorting options.
            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.
            cache_dir: optional directory to cache product metadata XML files for
                harmonization. If not set, XML is fetched each time (no caching).
            context: the data source context.
        """
        self.harmonize = harmonize

        # Set up cache directory for product metadata XML.
        self._cache_dir: UPath | None = None
        if cache_dir is not None:
            if context.ds_path is not None:
                self._cache_dir = join_upath(context.ds_path, cache_dir)
            else:
                self._cache_dir = UPath(cache_dir)
            self._cache_dir.mkdir(parents=True, exist_ok=True)

        selected_band_sets = select_band_sets(
            self.BAND_SETS_NEEDING_HARMONIZATION
            + self.BAND_SETS_NOT_NEEDING_HARMONIZATION,
            band_sets,
            context,
        )

        super().__init__(
            collection=self.COLLECTION,
            band_sets=selected_band_sets,
            query=query,
            retry_config=retry_config,
            timeout=timeout,
            skip_items_missing_assets=skip_items_missing_assets,
            context=context,
        )

    def _get_product_xml(
        self, item_name: str, provider_name: str, provider: DataProviderItem
    ) -> Element | None:
        """Fetch the product metadata XML for a Sentinel-2 item from a provider.

        Looks for the product-metadata asset on the given provider.
        Results are cached on disk if cache_dir is configured.

        Args:
            item_name: the item name (used for caching).
            provider_name: the name of the data provider to get the XML from (used for caching).
            provider: the DataProviderItem.

        Returns:
            The parsed XML element, or None if not available.
        """
        # Check disk cache if configured.
        if self._cache_dir is not None:
            cache_path = self._cache_dir / f"{item_name}_{provider_name}.xml"
            if cache_path.exists():
                return ET.fromstring(cache_path.read_bytes())

        # Find the product-metadata asset on this provider.
        metadata_asset = None
        for asset in provider.assets:
            if asset.name == self.PRODUCT_METADATA_ASSET:
                metadata_asset = asset
                break

        if metadata_asset is None:
            return None

        # Fetch the XML.
        response = self._session.get(
            metadata_asset.url, timeout=self.timeout.total_seconds()
        )
        response.raise_for_status()
        xml_content = response.content
        tree = ET.fromstring(xml_content)

        # Cache to disk if configured.
        if self._cache_dir is not None:
            cache_path = self._cache_dir / f"{item_name}_{provider_name}.xml"
            cache_path.write_bytes(xml_content)

        return tree

    def _get_harmonize_offset_callback(
        self,
    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]]:
        """Get a callback that applies the standard harmonization offset.

        This is used when boa_offset_applied is True, indicating the offset was applied
        and needs to be removed for harmonization.

        Returns:
            A callback that subtracts the offset from the array.
        """
        offset = self.HARMONIZE_OFFSET

        def callback(array: npt.NDArray[Any]) -> npt.NDArray[Any]:
            # We expect to only apply the offset on single-band uint16 files.
            # Some files are not single-band/uint16 like TCI (3 uint8 bands), but those
            # exceptions don't need harmonization.
            if array.shape[0] != 1 or array.dtype != np.uint16:
                raise ValueError(
                    f"harmonization should only be applied to single-band uint16 arrays, but got {array.shape[0]}-band {array.dtype} array"
                )
            # We assume the offset is -1000 since that is the standard.
            # To avoid issues working with uint16 array, we clip to 1000+ and then subtract 1000.
            return np.clip(array, -offset, None) - (-offset)

        return callback

    @override
    def get_raster_callback(
        self,
        item_name: str,
        bands: list[str],
        provider_name: str,
        provider: DataProviderItem,
    ) -> Callable[[npt.NDArray[Any]], npt.NDArray[Any]] | None:
        """Get harmonization callback for a Sentinel-2 item if needed.

        Tries to fetch and parse the product metadata XML first, then falls back
        to checking the boa_offset_applied property.

        Args:
            item: the item being read.
            bands: the bands being read (TCI/visual doesn't need harmonization).
            provider_name: the name of the data provider being used.
            provider: the corresponding DataProviderItem.

        Returns:
            A harmonization callback, or None if not needed.
        """
        # TCI (visual) image does not need harmonization.
        if not self.harmonize or bands not in self.BAND_SETS_NEEDING_HARMONIZATION:
            return None

        # Try fetching product metadata XML first to determine harmonization.
        xml = self._get_product_xml(item_name, provider_name, provider)
        if xml is not None:
            return get_harmonize_callback(xml)

        # Fall back to boa_offset_applied property (AWS provider sets this explicitly).
        # - False: offset was NOT applied (older processing), no harmonization needed
        # - True: offset WAS applied (newer processing 04.00+), need to subtract it
        boa_offset_applied: bool | None = None
        if isinstance(provider.properties, Sentinel2DataProviderItemProperties):
            boa_offset_applied = provider.properties.boa_offset_applied
        if boa_offset_applied is False:
            return None
        if boa_offset_applied is True:
            return self._get_harmonize_offset_callback()

        # Could not determine harmonization status. Raise an error so the caller can
        # try another provider.
        raise ValueError(
            f"Could not determine harmonization status for item {item_name} from provider {provider_name}"
        )
