# 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

import collections
import logging
from datetime import date, datetime

import numpy
import xarray as xr
from dateutil.parser import parse
from odc.geo.geobox import GeoBox
from ows.wcs.v20 import ScaleAxis, ScaleExtent, ScaleSize, Slice, Trim
from rasterio import MemoryFile

from datacube_ows.loading import DataStacker
from datacube_ows.ogc_exceptions import WCS2Exception
from datacube_ows.ows_configuration import OWSConfig, get_config
from datacube_ows.resource_limits import ResourceLimited
from datacube_ows.utils import default_to_utc
from datacube_ows.wcs_scaler import WCSScaler, WCSScalerUnknownDimension

# from datacube_ows.wcs_utils import get_bands_from_styles

_LOG: logging.Logger = logging.getLogger(__name__)


def uniform_crs(cfg: OWSConfig, crs: str) -> str:
    """Helper function to transform a URL style EPSG definition to an 'EPSG:nnn' one"""
    if crs.startswith("http://www.opengis.net/def/crs/EPSG/"):
        code = crs.rpartition("/")[-1]
        crs = f"EPSG:{code}"
    elif crs.startswith("urn:ogc:def:crs:EPSG:"):
        code = crs.rpartition(":")[-1]
        crs = f"EPSG:{code}"
    elif crs.startswith("EPSG") or crs in cfg.published_CRSs:
        pass
    else:
        raise WCS2Exception(
            f"Not a CRS: {crs}",
            WCS2Exception.NOT_A_CRS,
            locator=crs,
            valid_keys=list(cfg.published_CRSs),
        )
    return crs


def get_coverage_data(request, styles, qprof) -> tuple:
    # pylint: disable=too-many-locals, protected-access

    cfg = get_config()
    qprof.start_event("setup")
    layer_name = request.coverage_id
    layer = cfg.layer_index.get(layer_name)
    if not layer or not layer.wcs:
        raise WCS2Exception(
            f"Invalid coverage: {layer_name}",
            WCS2Exception.NO_SUCH_COVERAGE,
            locator="COVERAGE parameter",
            valid_keys=list(cfg.layer_index),
        )

    #
    # CRS handling
    #

    native_crs = layer.native_CRS
    subsetting_crs = uniform_crs(cfg, request.subsetting_crs or native_crs)
    if subsetting_crs not in cfg.published_CRSs:
        raise WCS2Exception(
            f"Invalid subsettingCrs: {subsetting_crs}",
            WCS2Exception.SUBSETTING_CRS_NOT_SUPPORTED,
            locator=subsetting_crs,
            valid_keys=list(cfg.published_CRSs),
        )

    output_crs = uniform_crs(cfg, request.output_crs or subsetting_crs or native_crs)

    if output_crs not in cfg.published_CRSs:
        raise WCS2Exception(
            f"Invalid outputCrs: {output_crs}",
            WCS2Exception.OUTPUT_CRS_NOT_SUPPORTED,
            locator=output_crs,
            valid_keys=list(cfg.published_CRSs),
        )

    #
    # Subsetting/Scaling
    #

    scaler = WCSScaler(layer, subsetting_crs)
    times = layer.ranges.times

    subsets = request.subsets

    if len(subsets) != len({subset.dimension.lower() for subset in subsets}):
        dimensions = [subset.dimension.lower() for subset in subsets]
        duplicate_dimensions = [
            item for item, count in collections.Counter(dimensions).items() if count > 1
        ]

        raise WCS2Exception(
            f"Duplicate dimension{'s' if len(duplicate_dimensions) > 1 else ''}: "
            f"{', '.join(duplicate_dimensions)}",
            WCS2Exception.INVALID_SUBSETTING,
            locator=",".join(duplicate_dimensions),
        )

    for subset in subsets:
        dimension = subset.dimension.lower()
        if dimension == "time":
            if isinstance(subset, Trim):
                if "," in subset.high:
                    raise WCS2Exception(
                        "Subsets can only contain 2 elements - the lower and upper bounds. For arbitrary date lists, use WCS1",
                        WCS2Exception.INVALID_SUBSETTING,
                        locator="time",
                    )
                if layer.time_resolution.is_subday():
                    low: date | datetime | None = (
                        default_to_utc(parse(subset.low))
                        if subset.low is not None
                        else None
                    )
                    high: date | datetime | None = (
                        default_to_utc(parse(subset.high))
                        if subset.high is not None
                        else None
                    )
                else:
                    low = parse(subset.low).date() if subset.low is not None else None
                    high = (
                        parse(subset.high).date() if subset.high is not None else None
                    )
                if low is not None:
                    times = [time for time in times if time >= low]
                if high is not None:
                    times = [time for time in times if time <= high]
            elif isinstance(subset, Slice):
                point = parse(subset.point).date()
                times = [point]
        else:
            try:
                if isinstance(subset, Trim):
                    scaler.trim(dimension, subset.low, subset.high)
                elif isinstance(subset, Slice):
                    scaler.slice(dimension, subset.point)
            except WCSScalerUnknownDimension:
                raise WCS2Exception(
                    f"Invalid subsetting axis {subset.dimension}",
                    WCS2Exception.INVALID_AXIS_LABEL,
                    locator=subset.dimension,
                ) from None

    #
    # Transform spatial extent to native CRS.
    #
    scaler.to_crs(output_crs)

    #
    # Scaling
    #

    scales = request.scales
    if len(scales) != len({subset.axis.lower() for subset in scales}):
        axes = [subset.axis.lower() for subset in scales]
        duplicate_axes = [
            item for item, count in collections.Counter(axes).items() if count > 1
        ]
        raise WCS2Exception(
            f"Duplicate scales for ax{'i' if len(duplicate_axes) == 1 else 'e'}s: "
            f"{', '.join(duplicate_axes)}",
            WCS2Exception.INVALID_SCALE_FACTOR,
            locator=",".join(duplicate_axes),
        )

    for scale in scales:
        axis = scale.axis.lower()

        if axis in ("time", "k"):
            raise WCS2Exception(
                f"Cannot scale axis {scale.axis}",
                WCS2Exception.INVALID_SCALE_FACTOR,
                locator=scale.axis,
            )
        if isinstance(scale, ScaleAxis):
            scaler.scale_axis(axis, scale.factor)
        elif isinstance(scale, ScaleSize):
            scaler.scale_size(axis, scale.size)
        elif isinstance(scale, ScaleExtent):
            scaler.scale_extent(axis, scale.low, scale.high)

    #
    # Rangesubset
    #

    band_labels = layer.band_idx.band_labels(canonical=True)
    trans_band_labels = layer.band_idx.band_labels()
    if request.range_subset:
        bands = []
        for range_subset in request.range_subset:
            if isinstance(range_subset, str):
                if range_subset in band_labels:
                    bands.append(range_subset)
                elif range_subset in trans_band_labels:
                    bands.append(band_labels[trans_band_labels.index(range_subset)])
                else:
                    raise WCS2Exception(
                        f"No such field {range_subset}",
                        WCS2Exception.NO_SUCH_FIELD,
                        locator=range_subset,
                        valid_keys=band_labels,
                    )
            else:
                if (
                    range_subset.start not in band_labels
                    and range_subset.start not in trans_band_labels
                ):
                    raise WCS2Exception(
                        f"No such field {range_subset.start}",
                        WCS2Exception.ILLEGAL_FIELD_SEQUENCE,
                        locator=range_subset.start,
                        valid_keys=band_labels,
                    )
                if (
                    range_subset.end not in band_labels
                    and range_subset.end not in trans_band_labels
                ):
                    raise WCS2Exception(
                        f"No such field {range_subset.end}",
                        WCS2Exception.ILLEGAL_FIELD_SEQUENCE,
                        locator=range_subset.end,
                        valid_keys=band_labels,
                    )

                if range_subset.start in band_labels:
                    start = band_labels.index(range_subset.start)
                else:
                    start = trans_band_labels.index(range_subset.start)
                if range_subset.end in band_labels:
                    end = band_labels.index(range_subset.end)
                else:
                    end = trans_band_labels.index(range_subset.end)
                bands.extend(
                    band_labels[start : (end + 1) if end > start else (end - 1)]
                )
    # Uncomment to restore original styles parameter hack.
    #
    # elif styles:
    #     bands = get_bands_from_styles(styles, layer, version=2)
    #     if not bands:
    #         bands = band_labels
    else:
        bands = band_labels

    #
    # Format handling
    #

    if not request.format:
        fmt = cfg.wcs_formats_by_name[layer.native_format]  # type: ignore[index]
    else:
        try:
            fmt = cfg.wcs_formats_by_mime[request.format]
        except KeyError:
            raise WCS2Exception(
                f"Unsupported format: {request.format}",
                WCS2Exception.INVALID_PARAMETER_VALUE,
                locator="FORMAT",
                valid_keys=list(cfg.wcs_formats_by_mime),
            ) from None

    if len(times) > 1 and not fmt.multi_time:
        raise WCS2Exception(
            "Format does not support multi-time datasets - "
            "either constrain the time dimension or choose a different format",
            WCS2Exception.INVALID_SUBSETTING,
            locator="FORMAT or SUBSET",
        )
    affine = scaler.affine()
    geobox = GeoBox((scaler.size.y, scaler.size.x), affine, cfg.crs(output_crs))

    stacker = DataStacker(layer, geobox, times, bands=bands)
    qprof.end_event("setup")
    qprof.start_event("count-datasets")
    n_datasets = stacker.n_datasets()
    qprof.end_event("count-datasets")
    qprof["n_datasets"] = n_datasets

    try:
        layer.resource_limits.check_wcs(
            n_datasets,
            geobox.height,
            geobox.width,
            sum(layer.band_idx.dtype_size(b) for b in bands),
            len(times),
        )
    except ResourceLimited as e:
        if e.wcs_hard or not layer.low_res_product_names:
            raise WCS2Exception(
                f"This request processes too much data to be served in a reasonable amount of time. ({e}) "
                + "Please reduce the bounds of your request and try again."
            ) from None
        stacker.resource_limited = True
        qprof["resource_limited"] = str(e)

    if n_datasets == 0:
        raise WCS2Exception(
            "The requested spatio-temporal subsets return no data.",
            WCS2Exception.INVALID_SUBSETTING,
            http_response=404,
        )

    qprof.start_event("fetch-datasets")
    datasets = stacker.datasets()
    qprof.end_event("fetch-datasets")
    if qprof.active:
        qprof["datasets"] = {
            str(q): [str(i) for i in ids] for q, ids in stacker.dsids().items()
        }
    qprof.start_event("load-data")
    # FIXME: output can be None.
    output = stacker.data(datasets, skip_corrections=True)
    qprof.end_event("load-data")

    # Clean extent flag band from output
    raw_bands = [layer.band_idx.locale_band(b) for b in bands]
    for k, _ in output.data_vars.items():  # type: ignore[union-attr]
        if k not in raw_bands:
            output = output.drop_vars([k])  # type: ignore[union-attr]

    #
    # TODO: configurable
    #
    if fmt.mime == "image/geotiff":
        output = fmt.renderer(request.version)(
            request, output, output_crs, layer, scaler.size.x, scaler.size.y, affine
        )

    else:
        output = fmt.renderer(request.version)(request, output, output_crs)

    headers = {
        "Content-Type": fmt.mime,
        "content-disposition": f"attachment; filename={request.coverage_id}.{fmt.extension}",
    }
    headers.update(layer.resource_limits.wcs_cache_rules.cache_headers(n_datasets))
    return output, headers


def get_tiff(request, data, crs, product, width: int, height, affine):
    """Uses rasterio MemoryFiles in order to return a streamable GeoTiff response"""
    # Does not support multi-time dimension data - is this even possible in GeoTiff?
    supported_dtype_map = {
        "uint8": 1,
        "int8": 2,
        "uint16": 3,
        "int16": 4,
        "uint32": 5,
        "int32": 6,
        "float32": 7,
        "float64": 8,
        "complex": 10,
        "complex64": 11,
        "complex128": 12,
    }

    dtype_list = [data[array].dtype for array in data.data_vars]
    dtype = str(max(dtype_list, key=lambda d: supported_dtype_map[str(d)]))

    # TODO: convert other parameters as-well
    gtiff = request.geotiff_encoding_parameters
    cfg = get_config()

    if len(data.time) > 1:
        raise WCS2Exception("Multiple time slices not supported by GeoTIFF format")
    data = data.squeeze(dim="time", drop=True)
    data = data.astype(dtype)
    nodata = 0
    for band in data.data_vars:
        nodata = product.band_idx.nodata_val(band)
    with MemoryFile() as memfile:
        # pylint: disable=protected-access, bad-continuation

        kwargs = {}
        if gtiff.tile_width is not None:
            kwargs["blockxsize"] = gtiff.tile_width
        if gtiff.tile_height is not None:
            kwargs["blockysize"] = gtiff.tile_height

        if gtiff.predictor:
            predictor = gtiff.predictor.lower()
            if predictor == "horizontal":
                kwargs["predictor"] = 2
            elif predictor == "floatingpoint":
                kwargs["predictor"] = 3
        elif dtype == "float64":
            kwargs["predictor"] = 3
        else:
            kwargs["predictor"] = 2

        with memfile.open(
            driver="GTiff",
            width=width,
            height=height,
            count=len(data.data_vars),
            transform=affine,
            crs=crs,
            nodata=nodata,
            tiled=gtiff.tiling if gtiff.tiling is not None else True,
            compress=gtiff.compression.lower() if gtiff.compression else "lzw",
            interleave=gtiff.interleave or "band",
            dtype=dtype,
            **kwargs,
        ) as dst:
            for idx, band in enumerate(data.data_vars, start=1):
                dst.write(data[band].values, idx)
                dst.set_band_description(idx, product.band_idx.band_label(band))
                if cfg.wcs_tiff_statistics:
                    dst.update_tags(
                        idx, STATISTICS_MINIMUM=numpy.nanmin(data[band].values)
                    )
                    dst.update_tags(
                        idx, STATISTICS_MAXIMUM=numpy.nanmax(data[band].values)
                    )
                    dst.update_tags(
                        idx, STATISTICS_MEAN=numpy.nanmean(data[band].values)
                    )
                    dst.update_tags(
                        idx, STATISTICS_STDDEV=numpy.nanstd(data[band].values)
                    )
        return memfile.read()


def get_netcdf(request, data: xr.Dataset, crs) -> bytes:
    # Cleanup dataset attributes for NetCDF export
    data.attrs["crs"] = crs
    for v in data.data_vars.values():
        v.attrs["crs"] = crs
        if "spectral_definition" in v.attrs:
            del v.attrs["spectral_definition"]
        if "flags_definition" in v.attrs:
            del v.attrs["flags_definition"]
    if "time" in data and "units" in data["time"].attrs:
        del data["time"].attrs["units"]

    # And export to NetCDF
    return bytes(data.to_netcdf())
