# This file is part of datacube-ows, part of the Open Data Cube project.
# See https://opendatacube.org for more information.
#
# Copyright (c) 2017-2024 OWS Contributors
# SPDX-License-Identifier: Apache-2.0

from typing import cast

from typing_extensions import override
from xarray import DataArray, Dataset

from datacube_ows.config_utils import CFG_DICT, ConfigException
from datacube_ows.styles.base import StyleDefBase
from datacube_ows.styles.component import LINEAR_COMP_DICT, ComponentStyleDef
from datacube_ows.styles.ramp import ColorRampDef

TYPE_CHECKING = False
if TYPE_CHECKING:
    from datacube_ows.ows_configuration import OWSNamedLayer


class HybridStyleDef(ColorRampDef, ComponentStyleDef):
    """
    Hybrid component/colour ramp style type.

    Returns a linear blend of a component image and colour ramp image
    """

    auto_legend = False

    def __init__(
        self,
        product: "OWSNamedLayer",
        style_cfg: CFG_DICT,
        defer_multi_date: bool = False,
        stand_alone: bool = False,
        user_defined: bool = False,
    ) -> None:
        """
        See StyleBaseDef
        """
        super().__init__(
            product,
            style_cfg,
            defer_multi_date=defer_multi_date,
            stand_alone=stand_alone,
            user_defined=user_defined,
        )
        style_cfg = cast(CFG_DICT, self._raw_cfg)
        self.component_ratio = float(cast(float | str, style_cfg["component_ratio"]))
        if self.component_ratio < 0.0 or self.component_ratio > 1.0:
            raise ConfigException(
                "Component ratio must be a floating point number between 0 and 1"
            )

    @override
    def transform_single_date_data(self, data: Dataset) -> Dataset:
        """
        Apply style to raw data to make an RGBA image xarray (single time slice only)

        :param data: Raw data, all bands.
        :return: RGBA uint8 xarray
        """
        # pylint: disable=too-many-locals
        if self.index_function is not None:
            idx_data = self.index_function(data)
            data["index_function"] = (idx_data.dims, idx_data.data)

        imgdata = Dataset(coords=data)

        d: DataArray = data["index_function"]
        for band, _ in self.rgb_components.items():
            rampdata = DataArray(
                self.color_ramp.get_value(d, band), coords=d.coords, dims=d.dims
            )
            component_band_data: DataArray | None = None
            for c_band, c_intensity in cast(
                LINEAR_COMP_DICT, self.rgb_components[band]
            ).items():
                if callable(c_intensity):
                    imgband_component_data = cast(
                        DataArray, c_intensity(data[c_band], c_band, band)
                    )
                else:
                    imgband_component_data = data[c_band] * cast(DataArray, c_intensity)
                if component_band_data is not None:
                    component_band_data += imgband_component_data
                else:
                    component_band_data = imgband_component_data
                if band != "alpha":
                    component_band_data = self.compress_band(band, component_band_data)
            img_band_data = rampdata * 255.0 * (
                1.0 - self.component_ratio
            ) + self.component_ratio * cast(DataArray, component_band_data)
            imgdata[band] = (d.dims, img_band_data.astype("uint8").data)

        return imgdata


# Register HybridStyleDef as a priority subclass of StyleDefBase
#    (priority means takes precedence over ComponentStyleDef and ColourRampDef)
StyleDefBase.register_subclass(HybridStyleDef, "component_ratio", priority=True)
