# 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 numpy
import pytz
import xarray
from affine import Affine
from dateutil.parser import parse
from odc.geo import geom
from odc.geo.geobox import GeoBox
from ows.util import Version
from rasterio import MemoryFile

from datacube_ows.config_utils import ConfigException
from datacube_ows.loading import DataStacker
from datacube_ows.ogc_exceptions import WCS1Exception
from datacube_ows.ows_configuration import get_config
from datacube_ows.resource_limits import ResourceLimited
from datacube_ows.wcs_utils import get_bands_from_styles


class WCS1GetCoverageRequest:
    version = Version(1, 0, 0)
    # pylint: disable=too-many-instance-attributes, too-many-branches, too-many-statements, too-many-locals
    def __init__(self, args) -> None:
        self.args = args
        cfg = get_config()

        # Argument: Coverage (required)  -> product/layer
        if "coverage" not in args:
            raise WCS1Exception("No coverage specified",
                                WCS1Exception.MISSING_PARAMETER_VALUE,
                                locator="COVERAGE parameter",
                                valid_keys=list(cfg.layer_index))
        self.layer_name = args["coverage"]
        self.layer = cfg.layer_index.get(self.layer_name)
        if not self.layer or not self.layer.wcs:
            raise WCS1Exception("Invalid coverage: %s" % self.layer_name,
                                WCS1Exception.COVERAGE_NOT_DEFINED,
                                locator="COVERAGE parameter",
                                valid_keys=list(cfg.layer_index))

        # Argument: FORMAT (required) -> a supported format
        if "format" not in args:
            raise WCS1Exception("No FORMAT parameter supplied",
                                WCS1Exception.MISSING_PARAMETER_VALUE,
                                locator="FORMAT parameter",
                                valid_keys=cfg.wcs_formats_by_name)
        if args["format"] not in cfg.wcs_formats_by_name:
            raise WCS1Exception("Unsupported format: %s" % args["format"],
                                WCS1Exception.INVALID_PARAMETER_VALUE,
                                locator="FORMAT parameter",
                                valid_keys=cfg.wcs_formats_by_name)
        self.format = cfg.wcs_formats_by_name[args["format"]]

        # Argument: (request) CRS (required)
        if "crs" not in args:
            raise WCS1Exception("No request CRS specified",
                                WCS1Exception.MISSING_PARAMETER_VALUE,
                                locator="CRS parameter",
                                valid_keys=list(cfg.published_CRSs))
        self.request_crsid = args["crs"]
        if self.request_crsid not in cfg.published_CRSs:
            raise WCS1Exception("%s is not a supported CRS" % self.request_crsid,
                                WCS1Exception.INVALID_PARAMETER_VALUE,
                                locator="CRS parameter",
                                valid_keys=list(cfg.published_CRSs))
        self.request_crs = cfg.crs(self.request_crsid)

        # Argument: response_crs (optional)
        if "response_crs" in args:
            self.response_crsid = args["response_crs"]
            if self.response_crsid not in cfg.published_CRSs:
                raise WCS1Exception("%s is not a supported CRS" % self.response_crsid,
                                    WCS1Exception.INVALID_PARAMETER_VALUE,
                                    locator="RESPONSE_CRS parameter",
                                    valid_keys=list(cfg.published_CRSs))
            self.response_crs = cfg.crs(self.response_crsid)
        else:
            self.response_crsid = self.request_crsid
            self.response_crs = self.request_crs

        # Arguments: One of BBOX or TIME is required
        # if "bbox" not in args and "time" not in args:
        #    raise WCS1Exception("At least one of BBOX or TIME parameters must be supplied",
        #                        WCS1Exception.MISSING_PARAMETER_VALUE,
        #                        locator="BBOX or TIME parameter"
        #                        )

        # Argument: BBOX (technically not required if TIME supplied, but
        #       it's not clear to me what that would mean.)
        # For WCS 1.0.0 all bboxes will be specified as minx, miny, maxx, maxy
        if "bbox" not in args:
            raise WCS1Exception("No BBOX parameter supplied",
                                WCS1Exception.MISSING_PARAMETER_VALUE,
                                locator="BBOX or TIME parameter")
        try:
            self.minx, self.miny, self.maxx, self.maxy = map(float, args['bbox'].split(','))
        except:
            raise WCS1Exception("Invalid BBOX parameter",
                                WCS1Exception.INVALID_PARAMETER_VALUE,
                                locator="BBOX parameter")

        self.specified_search_extent = geom.polygon([(self.minx, self.miny),
                                        (self.minx, self.maxy),
                                        (self.maxx, self.maxy),
                                        (self.maxx, self.miny),
                                        (self.minx, self.miny)
                                       ],
                                       crs=self.request_crs
                                      )
        if self.request_crs == self.response_crs:
            self.extent = self.specified_search_extent
        else:
            # Convert to response_crs and rectify bounding box
            bbox = self.specified_search_extent.to_crs(self.response_crs).boundingbox
            self.minx = bbox.left
            self.maxx = bbox.right
            self.miny = bbox.bottom
            self.maxy = bbox.top
            self.extent = geom.polygon(
                [
                    (self.minx, self.miny), (self.minx, self.maxy),
                    (self.maxx, self.maxy), (self.maxx, self.miny),
                    (self.minx, self.miny)
               ],
               crs=self.response_crs
            )

        # Argument: TIME
        # if self.product.wcs_sole_time:
        #    self.times = [parse(self.product.wcs_sole_time).date()]
        if "time" not in args:
            #      CEOS treats no supplied time argument as all time.
            # I'm really not sure what the right thing to do is, but QGIS wants us to do SOMETHING - use configured
            # default.
            self.times = [self.layer.default_time]
        else:
            # TODO: the min/max/res format option?
            # It's a bit underspeced. I'm not sure what the "res" would look like.
            times = args["time"].split(",")
            self.times = []
            for t in times:
                if t == "now":
                    continue
                try:
                    time = parse(t).date()
                    if time not in self.layer.ranges.time_set:
                        raise WCS1Exception(
                            "Time value '%s' not a valid date for coverage %s" % (t, self.layer_name),
                            WCS1Exception.INVALID_PARAMETER_VALUE,
                            locator="TIME parameter",
                            valid_keys=[d.strftime('%Y-%m-%d') for d in self.layer.ranges.time_set]
                        )
                    self.times.append(time)
                except ValueError:
                    raise WCS1Exception(
                        "Time value '%s' not a valid ISO-8601 date" % t,
                        WCS1Exception.INVALID_PARAMETER_VALUE,
                        locator="TIME parameter",
                        valid_keys=[d.strftime('%Y-%m-%d') for d in self.layer.ranges.time_set]
                    )
            self.times.sort()

            if len(self.times) == 0:
                raise WCS1Exception(
                    "No valid ISO-8601 dates",
                    WCS1Exception.INVALID_PARAMETER_VALUE,
                    locator="TIME parameter",
                    valid_keys=[d.strftime('%Y-%m-%d') for d in self.layer.ranges.time_set]
                )
            elif len(self.times) > 1 and not self.format.multi_time:
                raise WCS1Exception(
                    f"Cannot select more than one time slice with the {self.format.name} format",
                    WCS1Exception.INVALID_PARAMETER_VALUE,
                    locator="TIME and FORMAT parameters"
                )

        # Range constraint parameter: MEASUREMENTS
        # No default is set in the DescribeCoverage, so it is required
        # But QGIS wants us to work without one, so take default from config
        if "measurements" in args:
            bands = args["measurements"]
            self.bands = []
            for b in bands.split(","):
                if not b:
                    continue
                try:
                    self.bands.append(self.layer.band_idx.locale_band(b))
                except ConfigException:
                    raise WCS1Exception(f"Invalid measurement: {b}",
                                        WCS1Exception.INVALID_PARAMETER_VALUE,
                                        locator="MEASUREMENTS parameter",
                                        valid_keys=self.layer.band_idx.band_labels())
            if not bands:
                raise WCS1Exception("No measurements supplied",
                                    WCS1Exception.INVALID_PARAMETER_VALUE,
                                    locator="MEASUREMENTS parameter",
                                    valid_keys = self.layer.band_idx.band_labels())
        elif "styles" in args and args["styles"]:
            # Use style bands.
            # Non-standard protocol extension.
            #
            # As we have correlated WCS and WMS service implementations,
            # we can accept a style from WMS, and return the bands used for it.
            self.bands = list(get_bands_from_styles(args["styles"], self.layer))
            if not self.bands:
                self.bands = self.layer.band_idx.band_labels()
        else:
            self.bands = self.layer.band_idx.band_labels()

        # Argument: EXCEPTIONS (optional - defaults to XML)
        if "exceptions" in args and args["exceptions"] != "application/vnd.ogc.se_xml":
            raise WCS1Exception(f"Unsupported exception format: {args['exceptions']}",
                                WCS1Exception.INVALID_PARAMETER_VALUE,
                                locator="EXCEPTIONS parameter")

        # Argument: INTERPOLATION (optional only nearest-neighbour currently supported.)
        #      If 'none' is supported in future, validation of width/height/res will need to change.
        if "interpolation" in args and args["interpolation"] != "nearest neighbor":
            raise WCS1Exception(f'Unsupported interpolation method: {args["interpolation"]}',
                                WCS1Exception.INVALID_PARAMETER_VALUE,
                                locator="INTERPOLATION parameter")

        if "width" in args:
            if "resx" in args or "resy" in args:
                raise WCS1Exception("Specify WIDTH/HEIGHT parameters OR RESX/RESY parameters - not both",
                                    WCS1Exception.MISSING_PARAMETER_VALUE,
                                    locator="RESX/RESY/WIDTH/HEIGHT parameters")
            if "height" not in args:
                raise WCS1Exception("WIDTH parameter supplied without HEIGHT parameter",
                                    WCS1Exception.MISSING_PARAMETER_VALUE,
                                    locator="WIDTH/HEIGHT parameters")
            try:
                self.height = int(args["height"])
                if self.height < 1:
                    raise ValueError()
            except ValueError:
                raise WCS1Exception("HEIGHT parameter must be a positive integer",
                                    WCS1Exception.INVALID_PARAMETER_VALUE,
                                    locator="HEIGHT parameter")
            try:
                self.width = int(args["width"])
                if self.width < 1:
                    raise ValueError()
            except ValueError:
                raise WCS1Exception("WIDTH parameter must be a positive integer",
                                    WCS1Exception.INVALID_PARAMETER_VALUE,
                                    locator="WIDTH parameter")

            self.resx = (self.maxx - self.minx) / self.width
            self.resy = (self.maxy - self.miny) / self.height
        elif "resx" in args:
            if "height" in args:
                raise WCS1Exception("Specify WIDTH/HEIGHT parameters OR RESX/RESY parameters - not both",
                                    WCS1Exception.MISSING_PARAMETER_VALUE,
                                    locator="RESX/RESY/WIDTH/HEIGHT parameters")
            if "resy" not in args:
                raise WCS1Exception("RESX parameter supplied without RESY parameter",
                                    WCS1Exception.MISSING_PARAMETER_VALUE,
                                    locator="RESX/RESY parameters")
            try:
                self.resx = float(args["resx"])
                if self.resx <= 0.0:
                    raise ValueError(0)
            except ValueError:
                raise WCS1Exception("RESX parameter must be a positive number",
                                    WCS1Exception.INVALID_PARAMETER_VALUE,
                                    locator="RESX parameter")
            try:
                self.resy = float(args["resy"])
                if self.resy <= 0.0:
                    raise ValueError(0)
            except ValueError:
                raise WCS1Exception("RESY parameter must be a positive number",
                                    WCS1Exception.INVALID_PARAMETER_VALUE,
                                    locator="RESY parameter")
            self.width = int(((self.maxx - self.minx) / self.resx) + 0.5)
            self.height = int(((self.maxy - self.miny) / self.resy) + 0.5)
        elif "height" in args:
            raise WCS1Exception("HEIGHT parameter supplied without WIDTH parameter",
                                WCS1Exception.MISSING_PARAMETER_VALUE,
                                locator="WIDTH/HEIGHT parameters")
        elif "resy" in args:
            raise WCS1Exception("RESY parameter supplied without RESX parameter",
                                WCS1Exception.MISSING_PARAMETER_VALUE,
                                locator="RESX/RESY parameters")
        else:
            raise WCS1Exception("You must specify either the WIDTH/HEIGHT parameters or RESX/RESY",
                                WCS1Exception.MISSING_PARAMETER_VALUE,
                                locator="RESX/RESY/WIDTH/HEIGHT parameters")

        xscale = (self.maxx - self.minx) / self.width
        yscale = (self.miny - self.maxy) / self.height
        trans_aff = Affine.translation(self.minx, self.maxy)
        scale_aff = Affine.scale(xscale, yscale)
        self.affine = trans_aff * scale_aff
        self.geobox = GeoBox((self.height, self.width), self.affine, self.response_crs)
        self.ows_stats = bool(args.get("ows_stats"))


def get_coverage_data(req, qprof) -> tuple | tuple[int, tuple]:
    # pylint: disable=too-many-locals, protected-access
    stacker = DataStacker(req.layer,
                          req.geobox,
                          req.times,
                          bands=req.bands)
    qprof.start_event("count-datasets")
    n_datasets = stacker.n_datasets()
    qprof.end_event("count-datasets")
    qprof["n_datasets"] = n_datasets

    try:
        req.layer.resource_limits.check_wcs(n_datasets,
                                            req.geobox.height, req.geobox.width,
                                            sum(req.layer.band_idx.dtype_size(b) for b in req.bands),
                                            len(req.times))
    except ResourceLimited as e:
        if e.wcs_hard or not req.layer.low_res_product_names:
            raise WCS1Exception(
                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.")
        stacker.resource_limited = True
        qprof["resource_limited"] = str(e)
    if n_datasets == 0:
        # Return an empty coverage file with full metadata?
        qprof.start_event("build_empty_dataset")
        cfg = get_config()
        x_range = (req.minx, req.maxx)
        y_range = (req.miny, req.maxy)
        xname = cfg.published_CRSs[req.response_crsid]["horizontal_coord"]
        yname = cfg.published_CRSs[req.response_crsid]["vertical_coord"]
        xvals = numpy.linspace(
            x_range[0],
            x_range[1],
            num=req.width
        )
        yvals = numpy.linspace(
            y_range[0],
            y_range[1],
            num=req.height
        )
        if req.layer.time_resolution.is_subday():
            timevals = [
                numpy.datetime64(dt.astimezone(pytz.utc).isoformat(), "ns")
                for dt in req.times
            ]
        else:
            timevals = req.times
        if cfg.published_CRSs[req.request_crsid]["vertical_coord_first"]:
            nparrays = {
                band: (("time", yname, xname),
                       numpy.full((len(req.times), len(yvals), len(xvals)),
                                  req.layer.band_idx.nodata_val(band),
                                  dtype=req.layer.band_idx.dtype_val(band))
                      )
                for band in req.bands
            }
        else:
            nparrays = {
                band: (("time", xname, yname),
                       numpy.full((len(req.times), len(xvals), len(yvals)),
                                  req.layer.band_idx.nodata_val(band),
                                  dtype=req.layer.band_idx.dtype_val(band))
                      )
                for band in req.bands
            }
        data = xarray.Dataset(
            nparrays,
            coords={
                "time": timevals,
                xname: xvals,
                yname: yvals,
            }
        ).astype("int16")
        qprof.start_event("end_empty_dataset")
        qprof["write_action"] = "Write Empty"

        return n_datasets, data

    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
    sanitised_bands = [req.layer.band_idx.locale_band(b) for b in req.bands]
    for k, v in output.data_vars.items():  # type: ignore[union-attr]
        if k not in sanitised_bands:
            output = output.drop_vars([k])  # type: ignore[union-attr]
    qprof["write_action"] = "Write Data"
    return n_datasets, output


def get_tiff(req, data):
    """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)]))

    data = data.squeeze(dim="time", drop=True)
    data = data.astype(dtype)
    cfg = get_config()
    xname = cfg.published_CRSs[req.response_crsid]["horizontal_coord"]
    yname = cfg.published_CRSs[req.response_crsid]["vertical_coord"]
    nodata = 0
    for band in data.data_vars:
        nodata = req.layer.band_idx.nodata_val(band)
    with MemoryFile() as memfile:
        # pylint: disable=protected-access, bad-continuation
        with memfile.open(
            driver="GTiff",
            width=data.sizes[xname],
            height=data.sizes[yname],
            count=len(data.data_vars),
            transform=req.affine,
            crs=req.response_crsid,
            nodata=nodata,
            tiled=True,
            compress="lzw",
            interleave="band",
            dtype=dtype) as dst:
            for idx, band in enumerate(data.data_vars, start=1):
                dst.write(data[band].values, idx)
                dst.set_band_description(idx, req.layer.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(req, data):
    # Cleanup dataset attributes for NetCDF export
    data.attrs["crs"] = req.response_crsid
    for k, v in data.data_vars.items():
        v.attrs["crs"] = req.response_crsid
        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 data.to_netcdf()
