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

import logging

logger = logging.getLogger(__name__)    

from contextlib import contextmanager 

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

    def __enter__(self):
        import contextlib, zarr

        #self._outer_object = self._exit_stack.enter_context(
        #self._outer_object = zarr.open(store=self.outer_path, mode=('r' if self.readonly else 'a'))
        self._outer_object = zarr.open(store=self.outer_path)

        #self._exit_stack = contextlib.ExitStack()        
        #self._exit_stack.enter_context(self._outer_object.store.fs.transaction)
            
        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()
        pass


    def ensure_group(self, name, attrs=None):
        # Makes sure `name` exists in the current zarr store.
        # Possibly updates attributes.
        node = self.outer_object(None)
        if name not in node:
            obj = node.create_group(current)
        else:
            obj = node[name]

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

        return obj


    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:
            path: name or path of the dataset        
            dtype: data type (numpy-dtype)
            frame_shape: shape of the data point; note that
              the actual shape of the resulting data set will contain
              more dimension
            attrs: attributes to set for the dataset object
        '''

        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="{str(base.store)}" name="{path}" '
                        f'shape="{dshape}" dtype={dtype}')
            dset = base.create_array(path,
                                     shape=dshape,
                                     #maxshape=(None, *dshape[1:]),
                                     dtype=dtype)

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

        return dset
