from camagick.storage.base import ArrayAccess as ArrayAccessBase
from camagick.storage.base import RefuseIndex, RefuseOverwrite, RefuseSparse

import numpy as np
import h5py, logging, os, fsspec

logger = logging.getLogger(__name__)    

from contextlib import contextmanager 

class ArrayAccess(ArrayAccessBase):
    '''
    Implementation of a HDF5 file array access.
    '''

    #def __init__(self, outer=None, mode=None, **tags):
    #    '''
    #    Initializes an HDF5 array storage access wrapper.
    #
    #    Args:
    #
    #        outer: The outer part of the storage (i.e. path of the HDF file)
    #
    #        mode: Mode, as specified by the `ArrayAccess` base API. Note that
    #          this is _not_ the `h5py.File` mode!
    #
    #        tags: key=value access list, where "key" is the local
    #          handle, and "value" is the corresponding inner path.
    #          The value specified here is in addition to the "sdir"
    #          part (right?)
    #    '''
    #
    #    # HDF5 needs to have the root node named explicitly if
    #    # we want to access attributes like .keys or .attrs
    #    super().__init__(outer, mode, **tags)


    def _recurse_obj_fname(self, x):
        if hasattr(x, "filename"):
            return x.filename
        elif hasattr(x, "file") and x.file not in (x, None):
            return self._recurse_obj_fname(x.file)
        elif hasattr(x, "parent") and x.parent not in (x, None):
            return self._recurse_obj_fname(x.parent)
        return ''


    def __enter__(self):
        import contextlib
        self._exit_stack = contextlib.ExitStack()
        
        if self.readonly:
            self._file_obj = self._exit_stack.enter_context(
                fsspec.open(self.outer_path, mode='rb'))
            self._outer_object = self._exit_stack.enter_context(
                h5py.File(self._file_obj, 'r'))
            
        else:
            self._file_obj = self._exit_stack.enter_context(
                fsspec.open(self.outer_path, mode='rb+'))
            parent_dir = self._file_obj.fs._parent(self._file_obj.path)
            if not self._file_obj.fs.exists(parent_dir):
                self._file_obj.fs.mkdir(parent_dir)
            self._outer_object = self._exit_stack.enter_context(
                h5py.File(self._file_obj, 'a'))
            
        return self


    def outer_object(self, tag):
        # Ignore the tag -- this backend currently only supports single backend
        return self._outer_object
    

    def __exit__(self, *args):
        self._exit_stack.close()


    def ensure_group(self, name, attrs=None):
        # Makes sure `nodepath` exists in the current h5 file and is a data group
        nodepath = name
        subpath = ""
        subnode = self.outer_object(None)
        for current in nodepath.split('/'):

            if len(current) == 0:
                continue # nothing to do
            try:
                subnode = subnode[current]
                
            except KeyError:
                logger.info(f'msg="Creating node" file="{self._recurse_obj_fname(subnode)}" '
                            f'parent="{subnode.name}" '
                            f'name="{current}"')
                subnode = subnode.create_group(current)

            if attrs is not None:
                subnode.attrs.update(attrs)

        return subnode


    def _dtype_spec_to_dtype(self, ds):
        if isinstance(ds, str):
            try:            
                return {
                    'number': float,
                    'integer': int,
                    'array': float,
                }[ds]
            except KeyError:
                errmsg = f'tag="{key}" msg="Don\'t know how to save" dtype="{dt_spec}"'
                logger.error(errmsg)
                raise RuntimeError(errmsg)

        else:
            return ds


    def ensure_dataset(self, path, dtype=None, frame_shape=None, num_frames=0,
                       attrs=None):
        '''
        Received with the first data point, sets names, types etc.

        Args:
            name: name or path of the dataset        
            dtype: data type (numpy-dtype)
            point_shape: shape of the data point; note that
              the actual shape of the resulting data set will contain
              more dimension
        '''

        base = self.outer_object(None)

        if self.readonly:
            return base[path] if len(path) > 0 else base
        
        dshape = (0, *frame_shape) if frame_shape is not None else (0,)
        
        try:
            dset = base[path]

            if (dset.dtype != dtype):
                if  'O' not in self._frame_mode:
                    raise RuntimeError(f'msg="Dataset exists with different dtype"'
                                       f'have="{dset.dtype}" want="{dtype}"')
                else:
                    del base[path]
                    raise KeyError()

            if dset.shape[1:] != dshape[1:]:
                if  'O' not in self._frame_mode:                
                    raise RuntimeError(f'msg="Dataset exists with different shape"'
                                       f'have="{dset.shape}" want="{dshape}"')
                else:
                    del base[path]
                    raise KeyError()

        except KeyError:
            logger.info(f'msg="Creating dataset" file="{base.file.filename}" name="{path}" '
                        f'shape="{dshape}" dtype={dtype}')
            dset = base.create_dataset(path, dshape, maxshape=(None, *dshape[1:]),
                                       dtype=dtype, compression='lzf')

        if attrs is not None:
            dset.attrs.update(attrs)

        return dset
