# coding: utf-8

# ====================================================
# imports
from __future__ import annotations

import pickle
import numpy as np
from h5py import string_dtype
from numbers import Number

import numpy.typing as npt
from typing import Any
from typing import Mapping
from typing import TYPE_CHECKING

from ch5mpy.objects import File
from ch5mpy.objects import Group
from ch5mpy.utils import is_sequence

if TYPE_CHECKING:
    from ch5mpy import H5Array


# ====================================================
# code
def write_attribute(group: Group, name: str, obj: Any) -> None:
    """Write a simple object as a H5 group attribute."""
    try:
        group.attrs[name] = "None" if obj is None else obj

    except TypeError:
        # since <obj> cannot be stored directly, we pickle it
        # here we use numpy's void type to allow storing bytes generated by pickle
        group.attrs[name] = np.void(pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL))


def write_attributes(group: Group, **kwargs: Any) -> None:
    """Write multiple object as h5 group attributes."""
    for name, obj in kwargs.items():
        write_attribute(group, name, obj)


def _store_dataset(
    loc: Group | File, name: str, array: npt.NDArray[Any] | H5Array[Any]
) -> None:
    """Store a dataset."""
    if np.issubdtype(array.dtype, np.str_):
        loc.create_dataset(name, data=array.astype(object), dtype=string_dtype())

    else:
        loc.create_dataset(name, data=array)


def write_dataset(loc: Group | File, name: str, obj: Any) -> None:
    """Write an array-like object to a H5 dataset."""
    if isinstance(obj, Mapping):
        group = loc.create_group(name)
        write_datasets(group, **obj)

    else:
        # cast to np.array if needed (to get shape and dtype)
        array = np.array(obj) if not hasattr(obj, "shape") else obj

        if name in loc.keys():
            if loc[name] is array:
                # this exact dataset is already stored, do nothing
                return

            elif loc[name].shape == array.shape and loc[name].dtype == array.dtype:
                # a similar array already exists, simply copy the data
                loc[name][()] = array

            else:
                # a different array was stored, delete it before storing the new array
                del loc[name]
                _store_dataset(loc, name, array)

        else:
            _store_dataset(loc, name, array)


def write_datasets(loc: Group | File, **kwargs: Any) -> None:
    """Write multiple array-like objects to H5 datasets."""
    for name, obj in kwargs.items():
        write_dataset(loc, name, obj)


def write_object(loc: Group | File, name: str, obj: Any) -> None:
    """Write any object to a H5 file."""
    if isinstance(obj, Mapping):
        group = loc.create_group(name)
        write_objects(group, **obj)

    elif is_sequence(obj) or isinstance(obj, (Number, str)):
        write_dataset(loc, name, obj)

    else:
        raise NotImplementedError


def write_objects(loc: Group | File, **kwargs: Any) -> None:
    """Write multiple objects of any type to a H5 file."""
    for name, obj in kwargs.items():
        write_object(loc, name, obj)
