# The MIT License (MIT)
# Copyright (c) 2024 ESA Climate Change Initiative
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from typing import Any, Dict, Hashable, Mapping, Optional, Tuple

import dask.array as da
import numpy as np
import xarray
from xcube.core.store import (DataDescriptor, DataType, DataTypeLike,
                              VariableDescriptor)
from xcube.util.assertions import assert_true
from xcube.util.jsonschema import JsonIntegerSchema, JsonObjectSchema

from .constants import LOG


class VectorDataCube(xarray.Dataset):
    """A wrapper class around an xarray Dataset to keep it separate."""
    __slots__ = ()

VECTOR_DATA_CUBE_TYPE = DataType(
    VectorDataCube,
    ['vectordatacube', 'xarray.Dataset']
)

DataType.register_data_type(VECTOR_DATA_CUBE_TYPE)


def _attrs_to_json(attrs: Mapping[Hashable, Any]) \
        -> Optional[Dict[str, Any]]:
    new_attrs: Dict[str, Any] = {}
    for k, v in attrs.items():
        if isinstance(v, np.ndarray):
            v = v.tolist()
        elif isinstance(v, da.Array):
            v = np.array(v).tolist()
        if isinstance(v, float) and np.isnan(v):
            v = None
        new_attrs[str(k)] = v
    return new_attrs


class VectorDataCubeDescriptor(DataDescriptor):
    """
    A descriptor for an N-dimensional vector data cube whose spatial properties are
    designated by a geometry cooordinate. The cube is represented by an xarray.Dataset.
    Comprises a description of the data variables contained in the dataset.

    Regarding *time_range* and *time_period* parameters, please refer to
    https://github.com/dcs4cop/xcube/blob/master/docs/source/storeconv.md#date-time-and-duration-specifications

    :param data_id: An identifier for the data
    :param data_type: The data type of the data described
    :param crs: A coordinate reference system identifier,
        as an EPSG, PROJ or WKT string
    :param bbox: A bounding box of the data
    :param time_range: Start and end time delimiting this data's
        temporal extent
    :param time_period: The data's periodicity
        if it is evenly temporally resolved
    :param dims: A mapping of the dataset's dimensions to their sizes
    :param coords: mapping of the dataset's data coordinate names
        to instances of :class:VariableDescriptor
    :param data_vars: A mapping of the dataset's variable names
        to instances of :class:VariableDescriptor
    :param attrs: A mapping containing arbitrary attributes of the dataset
    :param open_params_schema: A JSON schema describing the parameters
        that may be used to open this data
    """

    def __init__(self,
                 data_id: str,
                 *,
                 data_type: DataTypeLike = VECTOR_DATA_CUBE_TYPE,
                 crs: str = None,
                 bbox: Tuple[float, float, float, float] = None,
                 time_range: Tuple[Optional[str], Optional[str]] = None,
                 time_period: str = None,
                 spatial_res: float = None,
                 dims: Mapping[str, int] = None,
                 coords: Mapping[str, 'VariableDescriptor'] = None,
                 data_vars: Mapping[str, 'VariableDescriptor'] = None,
                 attrs: Mapping[Hashable, any] = None,
                 open_params_schema: JsonObjectSchema = None,
                 **additional_properties):
        super().__init__(data_id=data_id,
                         data_type=data_type,
                         crs=crs,
                         bbox=bbox,
                         time_range=time_range,
                         time_period=time_period,
                         open_params_schema=open_params_schema)
        assert_true(VECTOR_DATA_CUBE_TYPE.is_super_type_of(data_type),
                    f'illegal data_type,'
                    f' must be compatible with {VECTOR_DATA_CUBE_TYPE!r}')
        if additional_properties:
            LOG.info(f'Additional properties received;'
                     f' will be ignored: {additional_properties}')
        self.dims = dict(dims) if dims else None
        self.spatial_res = spatial_res
        self.coords = coords if coords else None
        self.data_vars = data_vars if data_vars else None
        self.attrs = _attrs_to_json(attrs) if attrs else None

    @classmethod
    def get_schema(cls) -> JsonObjectSchema:
        schema = super().get_schema()
        schema.properties.update(
            dims=JsonObjectSchema(
                additional_properties=JsonIntegerSchema(minimum=0)
            ),
            coords=JsonObjectSchema(
                additional_properties=VariableDescriptor.get_schema()
            ),
            data_vars=JsonObjectSchema(
                additional_properties=VariableDescriptor.get_schema()
            ),
            attrs=JsonObjectSchema(
                additional_properties=True
            ),
        )
        schema.required = ['data_id', 'data_type']
        schema.additional_properties = False
        schema.factory = cls
        return schema
