from __future__ import annotations

import functools
import io
import os
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any

import numpy as np

from xarray.backends.common import (
    BACKEND_ENTRYPOINTS,
    BackendEntrypoint,
    BytesIOProxy,
    T_PathFileOrDataStore,
    WritableCFDataStore,
    _normalize_path,
    _open_remote_file,
    collect_ancestor_dimensions,
    datatree_from_dict_with_io_cleanup,
    find_root_and_group,
)
from xarray.backends.file_manager import (
    CachingFileManager,
    DummyFileManager,
    FileManager,
)
from xarray.backends.locks import HDF5_LOCK, combine_locks, ensure_lock, get_write_lock
from xarray.backends.netCDF4_ import (
    BaseNetCDF4Array,
    _build_and_get_enum,
    _encode_nc4_variable,
    _ensure_no_forward_slash_in_name,
    _extract_nc4_variable_encoding,
    _get_datatype,
    _nc4_require_group,
)
from xarray.backends.store import StoreBackendEntrypoint
from xarray.core import indexing
from xarray.core.utils import (
    FrozenDict,
    emit_user_level_warning,
    is_remote_uri,
    read_magic_number_from_file,
    try_read_magic_number_from_file_or_path,
)
from xarray.core.variable import Variable

if TYPE_CHECKING:
    import h5netcdf

    from xarray.backends.common import AbstractDataStore
    from xarray.core.dataset import Dataset
    from xarray.core.datatree import DataTree
    from xarray.core.types import ReadBuffer


class H5NetCDFArrayWrapper(BaseNetCDF4Array):
    def get_array(self, needs_lock=True):
        ds = self.datastore._acquire(needs_lock)
        return ds.variables[self.variable_name]

    def __getitem__(self, key):
        return indexing.explicit_indexing_adapter(
            key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
        )

    def _getitem(self, key):
        with self.datastore.lock:
            array = self.get_array(needs_lock=False)
            return array[key]


def _read_attributes(h5netcdf_var):
    # GH451
    # to ensure conventions decoding works properly on Python 3, decode all
    # bytes attributes to strings
    attrs = {}
    for k, v in h5netcdf_var.attrs.items():
        if k not in ["_FillValue", "missing_value"] and isinstance(v, bytes):
            try:
                v = v.decode("utf-8")
            except UnicodeDecodeError:
                emit_user_level_warning(
                    f"'utf-8' codec can't decode bytes for attribute "
                    f"{k!r} of h5netcdf object {h5netcdf_var.name!r}, "
                    f"returning bytes undecoded.",
                    UnicodeWarning,
                )
        attrs[k] = v
    return attrs


_extract_h5nc_encoding = functools.partial(
    _extract_nc4_variable_encoding,
    lsd_okay=False,
    h5py_okay=True,
    backend="h5netcdf",
    unlimited_dims=None,
)


def _h5netcdf_create_group(dataset, name):
    return dataset.create_group(name)


class H5NetCDFStore(WritableCFDataStore):
    """Store for reading and writing data via h5netcdf"""

    __slots__ = (
        "_filename",
        "_group",
        "_manager",
        "_mode",
        "autoclose",
        "format",
        "is_remote",
        "lock",
    )

    def __init__(
        self,
        manager: FileManager | h5netcdf.File | h5netcdf.Group,
        group=None,
        mode=None,
        lock=HDF5_LOCK,
        autoclose=False,
    ):
        import h5netcdf

        if isinstance(manager, h5netcdf.File | h5netcdf.Group):
            if group is None:
                root, group = find_root_and_group(manager)
            else:
                if type(manager) is not h5netcdf.File:
                    raise ValueError(
                        "must supply a h5netcdf.File if the group argument is provided"
                    )
                root = manager
            manager = DummyFileManager(root)

        self._manager = manager
        self._group = group
        self._mode = mode
        self.format = None
        # todo: utilizing find_root_and_group seems a bit clunky
        #  making filename available on h5netcdf.Group seems better
        self._filename = find_root_and_group(self.ds)[0].filename
        self.is_remote = is_remote_uri(self._filename)
        self.lock = ensure_lock(lock)
        self.autoclose = autoclose

    @classmethod
    def open(
        cls,
        filename,
        mode="r",
        format=None,
        group=None,
        lock=None,
        autoclose=False,
        invalid_netcdf=None,
        phony_dims=None,
        decode_vlen_strings=True,
        driver=None,
        driver_kwds=None,
        storage_options: dict[str, Any] | None = None,
    ):
        import h5netcdf

        if isinstance(filename, str) and is_remote_uri(filename) and driver is None:
            mode_ = "rb" if mode == "r" else mode
            filename = _open_remote_file(
                filename, mode=mode_, storage_options=storage_options
            )

        if isinstance(filename, BytesIOProxy):
            source = filename
            filename = io.BytesIO()
            source.getvalue = filename.getbuffer

        if isinstance(filename, io.IOBase) and mode == "r":
            magic_number = read_magic_number_from_file(filename)
            if not magic_number.startswith(b"\211HDF\r\n\032\n"):
                raise ValueError(
                    f"{magic_number!r} is not the signature of a valid netCDF4 file"
                )

        if format not in [None, "NETCDF4"]:
            raise ValueError("invalid format for h5netcdf backend")

        kwargs = {
            "invalid_netcdf": invalid_netcdf,
            "decode_vlen_strings": decode_vlen_strings,
            "driver": driver,
        }
        if driver_kwds is not None:
            kwargs.update(driver_kwds)
        if phony_dims is not None:
            kwargs["phony_dims"] = phony_dims

        if lock is None:
            if mode == "r":
                lock = HDF5_LOCK
            else:
                lock = combine_locks([HDF5_LOCK, get_write_lock(filename)])

        manager = (
            CachingFileManager(h5netcdf.File, filename, mode=mode, kwargs=kwargs)
            if isinstance(filename, str)
            else h5netcdf.File(filename, mode=mode, **kwargs)
        )
        return cls(manager, group=group, mode=mode, lock=lock, autoclose=autoclose)

    def _acquire(self, needs_lock=True):
        with self._manager.acquire_context(needs_lock) as root:
            ds = _nc4_require_group(
                root, self._group, self._mode, create_group=_h5netcdf_create_group
            )
        return ds

    @property
    def ds(self):
        return self._acquire()

    def open_store_variable(self, name, var):
        import h5netcdf
        import h5py

        dimensions = var.dimensions
        data = indexing.LazilyIndexedArray(H5NetCDFArrayWrapper(name, self))
        attrs = _read_attributes(var)

        # netCDF4 specific encoding
        encoding = {
            "chunksizes": var.chunks,
            "fletcher32": var.fletcher32,
            "shuffle": var.shuffle,
        }
        if var.chunks:
            encoding["preferred_chunks"] = dict(
                zip(var.dimensions, var.chunks, strict=True)
            )
        # Convert h5py-style compression options to NetCDF4-Python
        # style, if possible
        if var.compression == "gzip":
            encoding["zlib"] = True
            encoding["complevel"] = var.compression_opts
        elif var.compression is not None:
            encoding["compression"] = var.compression
            encoding["compression_opts"] = var.compression_opts

        # save source so __repr__ can detect if it's local or not
        encoding["source"] = self._filename
        encoding["original_shape"] = data.shape

        vlen_dtype = h5py.check_dtype(vlen=var.dtype)
        if vlen_dtype is str:
            encoding["dtype"] = str
        elif vlen_dtype is not None:  # pragma: no cover
            # xarray doesn't support writing arbitrary vlen dtypes yet.
            pass
        # just check if datatype is available and create dtype
        # this check can be removed if h5netcdf >= 1.4.0 for any environment
        elif (datatype := getattr(var, "datatype", None)) and isinstance(
            datatype, h5netcdf.core.EnumType
        ):
            encoding["dtype"] = np.dtype(
                data.dtype,
                metadata={
                    "enum": datatype.enum_dict,
                    "enum_name": datatype.name,
                },
            )
        else:
            encoding["dtype"] = var.dtype

        return Variable(dimensions, data, attrs, encoding)

    def get_variables(self):
        return FrozenDict(
            (k, self.open_store_variable(k, v)) for k, v in self.ds.variables.items()
        )

    def get_attrs(self):
        return FrozenDict(_read_attributes(self.ds))

    def get_dimensions(self):
        return FrozenDict((k, len(v)) for k, v in self.ds.dimensions.items())

    def get_parent_dimensions(self):
        return FrozenDict(collect_ancestor_dimensions(self.ds))

    def get_encoding(self):
        return {
            "unlimited_dims": {
                k for k, v in self.ds.dimensions.items() if v.isunlimited()
            }
        }

    def set_dimension(self, name, length, is_unlimited=False):
        _ensure_no_forward_slash_in_name(name)
        if is_unlimited:
            self.ds.dimensions[name] = None
            self.ds.resize_dimension(name, length)
        else:
            self.ds.dimensions[name] = length

    def set_attribute(self, key, value):
        self.ds.attrs[key] = value

    def encode_variable(self, variable, name=None):
        return _encode_nc4_variable(variable, name=name)

    def prepare_variable(
        self, name, variable, check_encoding=False, unlimited_dims=None
    ):
        import h5py

        _ensure_no_forward_slash_in_name(name)
        attrs = variable.attrs.copy()
        dtype = _get_datatype(variable, raise_on_invalid_encoding=check_encoding)

        fillvalue = attrs.pop("_FillValue", None)

        if dtype is str:
            dtype = h5py.special_dtype(vlen=str)

        # check enum metadata and use h5netcdf.core.EnumType
        if (
            hasattr(self.ds, "enumtypes")
            and (meta := np.dtype(dtype).metadata)
            and (e_name := meta.get("enum_name"))
            and (e_dict := meta.get("enum"))
        ):
            dtype = _build_and_get_enum(self, name, dtype, e_name, e_dict)
        encoding = _extract_h5nc_encoding(variable, raise_on_invalid=check_encoding)
        kwargs = {}

        # Convert from NetCDF4-Python style compression settings to h5py style
        # If both styles are used together, h5py takes precedence
        # If set_encoding=True, raise ValueError in case of mismatch
        if encoding.pop("zlib", False):
            if check_encoding and encoding.get("compression") not in (None, "gzip"):
                raise ValueError("'zlib' and 'compression' encodings mismatch")
            encoding.setdefault("compression", "gzip")

        if (
            check_encoding
            and "complevel" in encoding
            and "compression_opts" in encoding
            and encoding["complevel"] != encoding["compression_opts"]
        ):
            raise ValueError("'complevel' and 'compression_opts' encodings mismatch")
        complevel = encoding.pop("complevel", 0)
        if complevel != 0:
            encoding.setdefault("compression_opts", complevel)

        encoding["chunks"] = encoding.pop("chunksizes", None)

        # Do not apply compression, filters or chunking to scalars.
        if variable.shape:
            for key in [
                "compression",
                "compression_opts",
                "shuffle",
                "chunks",
                "fletcher32",
            ]:
                if key in encoding:
                    kwargs[key] = encoding[key]
        if name not in self.ds:
            nc4_var = self.ds.create_variable(
                name,
                dtype=dtype,
                dimensions=variable.dims,
                fillvalue=fillvalue,
                **kwargs,
            )
        else:
            nc4_var = self.ds[name]

        for k, v in attrs.items():
            nc4_var.attrs[k] = v

        target = H5NetCDFArrayWrapper(name, self)

        return target, variable.data

    def sync(self):
        self.ds.sync()

    def close(self, **kwargs):
        self._manager.close(**kwargs)


def _check_phony_dims(phony_dims):
    emit_phony_dims_warning = False
    if phony_dims is None:
        emit_phony_dims_warning = True
        phony_dims = "access"
    return emit_phony_dims_warning, phony_dims


def _emit_phony_dims_warning():
    emit_user_level_warning(
        "The 'phony_dims' kwarg now defaults to 'access'. "
        "Previously 'phony_dims=None' would raise an error. "
        "For full netcdf equivalence please use phony_dims='sort'.",
        UserWarning,
    )


def _normalize_filename_or_obj(
    filename_or_obj: T_PathFileOrDataStore,
) -> str | ReadBuffer | AbstractDataStore:
    if isinstance(filename_or_obj, bytes | memoryview):
        return io.BytesIO(filename_or_obj)
    else:
        return _normalize_path(filename_or_obj)


class H5netcdfBackendEntrypoint(BackendEntrypoint):
    """
    Backend for netCDF files based on the h5netcdf package.

    It can open ".nc", ".nc4", ".cdf" files but will only be
    selected as the default if the "netcdf4" engine is not available.

    Additionally it can open valid HDF5 files, see
    https://h5netcdf.org/#invalid-netcdf-files for more info.
    It will not be detected as valid backend for such files, so make
    sure to specify ``engine="h5netcdf"`` in ``open_dataset``.

    For more information about the underlying library, visit:
    https://h5netcdf.org

    See Also
    --------
    backends.H5NetCDFStore
    backends.NetCDF4BackendEntrypoint
    backends.ScipyBackendEntrypoint
    """

    description = (
        "Open netCDF (.nc, .nc4 and .cdf) and most HDF5 files using h5netcdf in Xarray"
    )
    url = "https://docs.xarray.dev/en/stable/generated/xarray.backends.H5netcdfBackendEntrypoint.html"

    def guess_can_open(self, filename_or_obj: T_PathFileOrDataStore) -> bool:
        filename_or_obj = _normalize_filename_or_obj(filename_or_obj)
        magic_number = try_read_magic_number_from_file_or_path(filename_or_obj)
        if magic_number is not None:
            return magic_number.startswith(b"\211HDF\r\n\032\n")

        if isinstance(filename_or_obj, str | os.PathLike):
            _, ext = os.path.splitext(filename_or_obj)
            return ext in {".nc", ".nc4", ".cdf"}

        return False

    def open_dataset(
        self,
        filename_or_obj: T_PathFileOrDataStore,
        *,
        mask_and_scale=True,
        decode_times=True,
        concat_characters=True,
        decode_coords=True,
        drop_variables: str | Iterable[str] | None = None,
        use_cftime=None,
        decode_timedelta=None,
        format=None,
        group=None,
        lock=None,
        invalid_netcdf=None,
        phony_dims=None,
        decode_vlen_strings=True,
        driver=None,
        driver_kwds=None,
        storage_options: dict[str, Any] | None = None,
    ) -> Dataset:
        # Keep this message for some versions
        # remove and set phony_dims="access" above
        emit_phony_dims_warning, phony_dims = _check_phony_dims(phony_dims)

        filename_or_obj = _normalize_filename_or_obj(filename_or_obj)
        store = H5NetCDFStore.open(
            filename_or_obj,
            format=format,
            group=group,
            lock=lock,
            invalid_netcdf=invalid_netcdf,
            phony_dims=phony_dims,
            decode_vlen_strings=decode_vlen_strings,
            driver=driver,
            driver_kwds=driver_kwds,
            storage_options=storage_options,
        )

        store_entrypoint = StoreBackendEntrypoint()

        ds = store_entrypoint.open_dataset(
            store,
            mask_and_scale=mask_and_scale,
            decode_times=decode_times,
            concat_characters=concat_characters,
            decode_coords=decode_coords,
            drop_variables=drop_variables,
            use_cftime=use_cftime,
            decode_timedelta=decode_timedelta,
        )

        # only warn if phony_dims exist in file
        # remove together with the above check
        # after some versions
        if store.ds._root._phony_dim_count > 0 and emit_phony_dims_warning:
            _emit_phony_dims_warning()

        return ds

    def open_datatree(
        self,
        filename_or_obj: T_PathFileOrDataStore,
        *,
        mask_and_scale=True,
        decode_times=True,
        concat_characters=True,
        decode_coords=True,
        drop_variables: str | Iterable[str] | None = None,
        use_cftime=None,
        decode_timedelta=None,
        format=None,
        group: str | None = None,
        lock=None,
        invalid_netcdf=None,
        phony_dims=None,
        decode_vlen_strings=True,
        driver=None,
        driver_kwds=None,
        **kwargs,
    ) -> DataTree:
        groups_dict = self.open_groups_as_dict(
            filename_or_obj,
            mask_and_scale=mask_and_scale,
            decode_times=decode_times,
            concat_characters=concat_characters,
            decode_coords=decode_coords,
            drop_variables=drop_variables,
            use_cftime=use_cftime,
            decode_timedelta=decode_timedelta,
            format=format,
            group=group,
            lock=lock,
            invalid_netcdf=invalid_netcdf,
            phony_dims=phony_dims,
            decode_vlen_strings=decode_vlen_strings,
            driver=driver,
            driver_kwds=driver_kwds,
            **kwargs,
        )

        return datatree_from_dict_with_io_cleanup(groups_dict)

    def open_groups_as_dict(
        self,
        filename_or_obj: T_PathFileOrDataStore,
        *,
        mask_and_scale=True,
        decode_times=True,
        concat_characters=True,
        decode_coords=True,
        drop_variables: str | Iterable[str] | None = None,
        use_cftime=None,
        decode_timedelta=None,
        format=None,
        group: str | None = None,
        lock=None,
        invalid_netcdf=None,
        phony_dims=None,
        decode_vlen_strings=True,
        driver=None,
        driver_kwds=None,
        **kwargs,
    ) -> dict[str, Dataset]:
        from xarray.backends.common import _iter_nc_groups
        from xarray.core.treenode import NodePath
        from xarray.core.utils import close_on_error

        # Keep this message for some versions
        # remove and set phony_dims="access" above
        emit_phony_dims_warning, phony_dims = _check_phony_dims(phony_dims)

        filename_or_obj = _normalize_filename_or_obj(filename_or_obj)
        store = H5NetCDFStore.open(
            filename_or_obj,
            format=format,
            group=group,
            lock=lock,
            invalid_netcdf=invalid_netcdf,
            phony_dims=phony_dims,
            decode_vlen_strings=decode_vlen_strings,
            driver=driver,
            driver_kwds=driver_kwds,
        )

        # Check for a group and make it a parent if it exists
        if group:
            parent = NodePath("/") / NodePath(group)
        else:
            parent = NodePath("/")

        manager = store._manager
        groups_dict = {}
        for path_group in _iter_nc_groups(store.ds, parent=parent):
            group_store = H5NetCDFStore(manager, group=path_group, **kwargs)
            store_entrypoint = StoreBackendEntrypoint()
            with close_on_error(group_store):
                group_ds = store_entrypoint.open_dataset(
                    group_store,
                    mask_and_scale=mask_and_scale,
                    decode_times=decode_times,
                    concat_characters=concat_characters,
                    decode_coords=decode_coords,
                    drop_variables=drop_variables,
                    use_cftime=use_cftime,
                    decode_timedelta=decode_timedelta,
                )

            if group:
                group_name = str(NodePath(path_group).relative_to(parent))
            else:
                group_name = str(NodePath(path_group))
            groups_dict[group_name] = group_ds

        # only warn if phony_dims exist in file
        # remove together with the above check
        # after some versions
        if store.ds._phony_dim_count > 0 and emit_phony_dims_warning:
            _emit_phony_dims_warning()

        return groups_dict


BACKEND_ENTRYPOINTS["h5netcdf"] = ("h5netcdf", H5netcdfBackendEntrypoint)