"""Utils to save and load h5 files."""

import os
from typing import Literal, Optional, Set, Union

import h5py
import numpy as np

from ..errors import FileLockedError
from ..types import ClassWithAsdict
from .data_transformation import transform_not_dict_on_save, transform_on_open


class LockFile:
    """Context manager that creates a lock file during writing to prevent conflicts.

    Usage:
        To use the LockFile context manager, simply wrap your code that writes to a file with it.
        This will ensure that only one process can write to the file at a time.

        Example:
        ```
        with LockFile("data.txt"):
             # Code that writes to the file
             with open("data.txt", "w") as f:
                 pass
        ```

    """

    def __init__(self, filename):
        """Initialize the LockFile object.

        Args:
            filename (str): The filename or filepath. Can end with an extension.
        """
        self.lock_filename = os.path.splitext(filename)[0] + ".lock"

    def __enter__(self):
        """Enter the context and create the lock file."""
        if os.path.exists(self.lock_filename):
            raise FileLockedError("File locked and cannot be opened in write mode")
        with open(self.lock_filename, "w", encoding="utf-8"):
            pass

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit the context and remove the lock file."""
        if os.path.exists(self.lock_filename):
            os.remove(self.lock_filename)


# -------------- Open file ----------------


def open_h5_group(
    group: Union[h5py.File, h5py.Group],
    key: Optional[Union[str, Set[str]]] = None,
    key_prefix: Optional[str] = None,
) -> dict:
    """Open h5 group and return dict.

    Args:
        group (Union[h5py.File, h5py.Group]): group to open
        key (Optional[Union[str, Set[str]]], optional): Keys to load. Defaults to None.
        key_prefix (Optional[str], optional): Key prefix. Defaults to None.

    Returns:
        dict: The loaded data. Each value was transform using transform_on_open function.
    """
    data = {}
    if key is not None:
        key = key if isinstance(key, set) else set([key])

    for group_key in group.keys():
        key = key if key_prefix is None else f"{key_prefix}/{key}"
        if key is not None and group_key not in key:
            continue
        value = group[group_key]
        if isinstance(value, h5py.Group):
            data[group_key] = open_h5_group(value)
        else:
            data[group_key] = transform_on_open(value[()])  # type: ignore
    return data


def open_h5(
    fullpath: str, key: Optional[Union[str, Set[str]]] = None, key_prefix: Optional[str] = None
) -> dict:
    """Open h5 file and return dict.

    Args:
        fullpath (str): Full filepath to the file.
        key (str | set[str], optional): Key or keys to load. Defaults to all keys.
        key_prefix (str, optional): Key prefix to put before each key. Defaults to None.

    Returns:
        dict: The loaded data. Each value was transform using transform_on_open function.
    """
    with h5py.File(fullpath, "r") as file:
        if key_prefix is None:
            return open_h5_group(file, key=key)
        return open_h5_group(file[key_prefix], key=key)  # type: ignore


# -------------- Save file ----------------


def save_sub_dict(
    group: Union[h5py.File, h5py.Group],
    data: Union[dict, list, np.ndarray, ClassWithAsdict],
    key: str,
    use_compression: Optional[Union[Literal[True], str]] = None,
):
    """Save the dict to a group. Each object is converted to a dict, array or simple value.

    Args:
        group (Union[h5py.File, h5py.Group]): h5 group where to save.
        data (Union[dict, list, np.ndarray, ClassWithAsdict]): Any object that either has 'asdict'
         or 'asarray' method, either is a dict, list or np.ndarray, either can be transformed
         using transform_not_dict_on_save function.
        key (str): Desire key for the data inside h5 file .
        use_compression (str|True, optional): If compression should be used.
         If true, 'gzip' is used, otherwise you can specify compression by providing a str.
         Defaults to not compressed.
    """
    if hasattr(data, "asdict"):
        data = data.asdict()  # type: ignore
    if hasattr(data, "asarray"):
        data = data.asarray()  # type: ignore
    if isinstance(data, dict):
        g = group.create_group(key)
        for k, v in data.items():
            save_sub_dict(g, v, k, use_compression=use_compression)
    elif (key is not None) and (data is not None):
        data = transform_not_dict_on_save(data)  # type: ignore
        if isinstance(data, (np.ndarray, list)):
            use_compression = "gzip" if use_compression is True else use_compression
            group.create_dataset(key, data=data, compression=use_compression)  # compression="gzip"
        else:
            group.create_dataset(key, data=data)


def save_dict(
    filename: str,
    data: dict,
    key_prefix: Optional[str] = None,
    use_compression: Optional[Union[Literal[True], str]] = None,
) -> float:
    """Save dict to h5 file.

    Args:
        filename (str): Full filepath to the file.
        data (dict): Data to save.
        key_prefix (str, optional): Key prefix of the desired location inside h5 file.
         If nested use `key1/key2`. Defaults to None, i.e. root of the file.
        use_compression (str|True, optional): If compression should be used. If true,
         'gzip' is used, otherwise you can specify compression by providing a str.
         Defaults to not compressed.

    Returns:
        float: Time of the last modification of the file.

    Example:
        ```
        save_dict('/path/to/file.h5', {'key1': 1})
        ```



    """
    dirname = os.path.dirname(filename)
    if dirname and not os.path.exists(dirname):
        os.makedirs(dirname, exist_ok=True)

    mode = "a"  # if os.path.exists(filename) else 'w'
    with LockFile(filename), h5py.File(filename, mode) as file:
        for key, value in data.items():
            key = key if key_prefix is None else f"{key_prefix}/{key}"
            if key in file:  # .keys():
                file.pop(key)
            if value is None:
                continue
            save_sub_dict(file, value, key, use_compression=use_compression)
    return os.path.getmtime(filename)


# -------------- Load keys ----------------


# /doc add example to this doc
def keys_h5(filename, key_prefix: Optional[str] = None) -> Set[str]:
    """
    Return the keys of an h5 file.

    Args:
        filename (str): Full filepath to the h5 file.
        key_prefix (str, optional): Key prefix to look at. If nested, use `key1/key2`.
            Defaults to None.

    Returns:
        Set[str]: Set of the keys in the h5 file.

    Example:
        ```
        >>> keys_h5('/path/to/file.h5', 'group1')
        {'key1', 'key2', 'key3'}
        ```
    """
    with h5py.File(filename, "r") as file:
        if key_prefix is not None:
            file = file[key_prefix]
        return set(file.keys())  # type: ignore


# -------------- Delete keys ----------------


def del_dict(
    filename: str,
    key: str,
    key_prefix: Optional[str] = None,
) -> float:
    """
    Delete a key from an h5 file.

    Args:
        filename (str): Full filepath to the file.
        key (str): The key to delete. If nested, use `key1/key2`.
        key_prefix (str, optional): If provided, the key prefix to use. Defaults to None.

    Returns:
        float: The time of the last modification of the file.

    Example:
        ```
        >>> del_dict('/path/to/file.h5', 'key1', 'prefix')
        "some_value"
        ```
    """
    with LockFile(filename), h5py.File(filename, "a") as file:
        key = key if key_prefix is None else f"{key_prefix}/{key}"
        file.pop(key)
    return os.path.getmtime(filename)
