'''
Array Storage API abstraction

This is meant to implement the operations that CAmagick supports in
an backend-agnostic way. Typically, CAmagick writes to HDF5, Zarr
and Tiled. So this API intents to support the common denominator of
all of these:

- Eternal / internal path (i.e. an entry point, "file" or "url"
  that designates a high-level resource, and an internal path
  that designates a low-level resource within the high-level
  one)

- Groups or Datasets (HDF5-like hiearchical structure, where
  there are datasets or groups, with the groups containing further
  datasets or subgroups, accessible by ["...path..."])

- Datasets can be sliced, and meta information (in particular data type
  and shape) can be accessed by `.dtype`, respectively `.shape`

- Datasets and/or groups can have attributes by means of `.attrs[...]`

- Files can be opened for reading, or for read-writing, optionally
  overwriting or refusing to create new files when opening for
  read-writing and the file already exists / doesn't exist.
'''

import time, logging
from numpy import array as np_array

logger = logging.getLogger(__name__)

__all__ = [
    "RefuseIndex",
    "RefuseOverwrite",
    "RefuseSparse",
    "ArrayAccess",
    "MemoryArrayAccess"
]

# list of procedures that return True if a specific storage
# path/url is of a specific format.

_format_probe = {
    'hdf5': lambda x: np_array([
        x.lower().endswith(s) for s in \
        ('.h5', '.hdf', '.hdf5', '.nx', '.cdf4')
    ]).any(),
    
    'zarr': lambda x: np_array([
        x.lower().endswith(s) for s in \
        ('.z', '.zarr')
    ]).any(),
}

class UnknownArrayType(RuntimeError): pass


def probe_url(url, backtrack=True):
    '''
    Tries to guess the "kind" of a data URL.

    Returns a named tuple (kind, outer, residual), where "kind" is a
    string that represents the type and is a submodule in
    `camagick.storage`, and "outer" is the outer-part of the URL
    that can be passed
    to any ArrayAccess object as a first parameter. "residual" is
    the remainder of the URL that didn't fit in "outer". This is
    in particular the case if the URL doesn't end with the
    storage format (e.g. "file:///path/to/file.hdf5/group/subgroup/").

    Probing formats is tricky business. There are 3 mechanisms in use,
    in this order:

      - Check the "scheme" part of the for something like "http+hdf5",
        where the part after "+" is the kind.

      - Check the extention of the final path element for known
        extensions, as in: "/path/to/file.hdf5", or possibly even
        "file:///path/to/file.hdf5#subgroup"

      - Backtracking the path elements one by one and check their
        extension, e.g. to find things like
        "/path/to/folder.zarr/subgroup/". This is only performed if
        `backtrack` is set to `True`.
    '''

    from urllib.parse import urlparse, urlunparse, ParseResult
    from collections import namedtuple

    purl = urlparse(url)

    UrlSpec = namedtuple('UrlSpec', ['kind', 'outer', 'residual'])

    # Step 1: check scheme (overrides any path element extensions)
    if purl.scheme not in (None, ''):
        tmp = purl.scheme.split('+')
        if len(tmp) == 2:
            return UrlSpec(tmp[1], url, '')

    # Steps 2+3: iteratively check file extensions of path elements
    path_elems = purl.path.split('/')
    nr_elems = len(path_elems)
    for j,elem in enumerate(path_elems[::-1]):
        i = nr_elems - j
        for kind, stest in _format_probe.items():
            if stest(elem):
                outer = urlunparse(ParseResult(
                    scheme=purl.scheme,
                    netloc=purl.netloc,
                    path='/'.join(path_elems[:i]),
                    params=purl.params,
                    query=purl.query,
                    fragment=None
                ))
                inner = '/'.join(path_elems[i:]) + \
                    (f'#{purl.fragment}' if purl.fragment not in (None, '') else '')
                # fragment=purl.fragment
                return UrlSpec(kind, outer, inner)

        if not backtrack:
            break

    raise UnknownArrayType(f'Unknown array type for: {url}')


def probe_kind(url):
    '''
    Wrapper for `probe_url_spec()` for backwards compatibility.
    Returns just the kind.
    '''
    return probe_url(url)[0]


def find_loader(url=None, kind=None):
    from importlib import import_module
    if kind is None:
        kind = probe_kind(url)
    mod = import_module(f'camagick.storage.{kind}')
    return getattr(mod, 'ArrayAccess')


from contextlib import contextmanager

@contextmanager
def load_from_url(url, kind=None, mode=None, **tags):
    '''
    Wrapper for find_loader(...) and the corresponding loading.
    '''
    Loader = find_loader(url, kind=kind)
    with Loader(url, mode, **tags) as obj:
        try:
            yield obj
        finally:
            pass


class RefuseFrameStorage(RuntimeError): pass

class RefuseShape(RefuseFrameStorage): pass

class RefuseIndex(RefuseFrameStorage): pass

class RefuseOverwrite(RefuseIndex): pass

class RefuseSparse(RefuseIndex): pass

class DatasetUnavailable(RuntimeError):
    def __init__(self, tag, reason=None):
        self.tag = tag
        super().__init__(f'Dataset unavailable: "{tag}" ({reason})')

class ArrayAccess:
    '''
    This is just to document the abstract API.

    The API elements are heavily tailored towards CAmagick's needs, which
    are essentially appending new data ("frames") to existing datasets.
    To implement a new backend, at least the following members need
    to be defined:

    - Store opening/closing: within context management infrastructure
      (`__enter__`, `__exit__`, and provide an `.outer_object` property)
    
    - Dataset management: `_ensure_dataset()`

    - (Optionally) dataset resizing: `._resize_dataset()` (and possibly
      `_reset_dataset()`), if the dataset doesn't have a `.resize(shape)`
      function.
  
    '''

    def __init__(self, outer=None, mode=None, **tags):
        '''
        Initializes the storage mechanism for "outer".

        We refer to resources in an hiearchical array storage by
        their URL. For every URL we differentiate between an
        "outer" part (prefix), pointing to the externally visible
        resource (e.g. HDF5 file), and an "inner" part, pointing
        to the resource within that file. Since we're focused on
        using a URL-like pattern, we're recycling the URL separator
        '#' for delimining outer from inner, so a typical URL might
        look like:
           "scheme://possibly.a.host.part/path/to/container#path/within/container"

        Here, everything up to '#' is the outer part, and everything
        after '#' is the inner part.

        One specific `ArrayAccess` implementation only covers one
        specific "outer" point -- this is sometimes important when
        the storage backend possibly locks, e.g. in the case of (writing)
        HDF5 files. This is the reason why we separate.

        We're liberal as to whether the repo URL is file#sub, or file/sub
        (maybe a little _too_ liberal...). As a consequence, here inner
        might be something like: "#...", or "/...", or just plain "...".
        But different backends react somewhat sensitive to that.
        We therefore give every backend the possibility to implement their
        own inner-path mangling by overriding the `._mangle_inner()` function.
        Mostly it's about detecting a '#', so this is the default implementation,
        but sometimes also checking whether there's an empty string involved
        (which HDF5 cannot handle) is important.
        
        Args:
            outer: Top-level store path (file name or URL).

            mode: What to do when we're required to create or overwrite a
              data point. We only concert ourselves with data points since
              this is what CAmagick's basic data currency is. For everything
              else (files, groups), the policy follows whatever must be done
              to ensure the data point mode policy can be followed.
        
              This must be either `None`, or a string consistinc or a
              combination of `a` (append), `o` or `O` (overwrite)
              or `x` (reset), possibly in combination with `+` (top-up),
              meaning:
        
                 - `a`: require that the data point be always appended
                   to the current dataset -- never overwrite or delete
                   existing data.
        
                 - `o`: overwrite already existing data points (frames).
                   If the shape of the new data point frames doesn't match
                   the old ones, then fail.

                 - `O`: same as `o`, but additioally delete and re-create
                   whole dataset entries, if the shape or data type of the
                   existing dataset does not match the new request.

                 - `x`: reset the entire dataset (i.e. delete all existing
                   data points) if a data point already exists at the
                   specified index.

                 - `+`: top-up marker will fill up empty slots, but without
                    actually writing any data to them. They will remain
                    filled with default settings of the corresponding backend
                    and the data type. If top-up is not specified, otherwise
                    the operation will fail if index is higher than the one of
                    the next slot in turn.

              If `mode` is set to `None`, each dataset, and consequently the
              whole store, is opened for reading only. This also reflects
              the behavior of the async environment enter/exit functionality.

            **tags: Keys are tag names (by which to refer the data in
              `.push_frame()`, and values are path elements, for each
              frame in itself.
        '''
        self._stor_spec = outer
        self._stor_dspath = {
            k:self._mangle_inner(v) for k,v in tags.items()
        }

        self._frame_mode = mode

        
    def _mangle_inner(self, inner):
        '''
        Mangle the inner path for an array storage resource (group or dataset).
        '''

        if len(inner) == 0:
            return '/'

        if inner[0] == '#':
            if len(inner) > 1:
                inner = inner[1:]
            else:
                inner = ''

        if len(inner) == 0:
            inner = '/'

        return inner


    def outer_obj(self, tag):
        raise RuntimeError('not implemented')


    @property
    def outer_path(self):
        '''
        Returns the "outer" part of the store specification.
        '''
        return self._stor_spec


    #@property
    #def inner_path(self):
    #    '''
    #    Returns the "inner" part of the store specification.
    #    '''
    #    return self._stor_spec[1]


    @property
    def readonly(self):
        return self._frame_mode is None


    # async wrappers for context management, which might block
    async def __aenter__(self):
        import concurrent, asyncio
        loop = asyncio.get_running_loop()
        with concurrent.futures.ThreadPoolExecutor() as pool:
            return await loop.run_in_executor(pool, self.__enter__)
        

    async def __aexit__(self, *args):
        import concurrent, asyncio
        loop = asyncio.get_running_loop()
        with concurrent.futures.ThreadPoolExecutor() as pool:
            return await loop.run_in_executor(pool, self.__exit__)

            
    def _split_dset_path(self, dset_path, splitter='/'):
        # Helper that will return the (node path, dataset path)
        # out of a whole path.
        # If there's only one component, assume it's a dataset.
        parts = dset_path.split(splitter)
        
        if len(parts) == 0:
            raise RuntimeError(f'invalid dataset path {dset_path}')

        if len(parts) == 1:
            return ('', parts[0])

        return ('/'.join(parts[:-1]), parts[-1])


    def ensure_group(self, name, attrs=None):
        '''
        Makse sure the group exists and has the proper attributes

        Args:
            name: The absolute group name within `outer_object`
        
            attrs: additional group attributes to set

        Returns a node object for the group.
        '''
        raise RuntimeError('not implemented')


    def ensure_dataset(self, path, dtype=None, frame_shape=None, num_frames=0,
                       attrs=None):
        '''
        Ensure the existence of the target dataset.
        
        This is typically called with the first received frame-push (i.e.
        data point), and is intended to set sets names, types,
        attributes etc.

        Args:
            path: full path of the dataset inside `base`
                
            dtype: data type (numpy-dtype); if array is in read-only
              mode, this and `frame_shape` are expected to be `None`,
              and the actual data type and frame shape of the target
              object will be ignored.
        
            frame_shape: shape of the data point; note that
              the actual shape of the resulting data set will contain
              one extra dimension upfront (i.e. the frame index).

            num_frames: Initial number of frames (i.e. the number of entries
             in the first dimension). `If this is 0, a dataset with 0 frames
             initially will be created, requiring `.push_frame()` to add data.
        
            *attrs: additional dataset attributes to set

        Return the dataset object.
        '''
        raise RuntimeError('not implemented')


    def _resize_dataset(self, dset, new_shape):
        '''
        Helper for `._ensure_frame()`.

        Essentially, this is the implementation-dependent part. Most backends
        will simply offer a `.resize()` call (this is true for HDF5 and Zarr),
        which is also the default implementation here.

        But others may require a more elaborate procedure (e.g. saving data,
        resizing, restoring).
        '''
        dset.resize(new_shape)


    def _reset_dataset(self, dset):
        '''
        Helper for `._ensure_frame()`.

        Erases all data points from a dataset.
        '''
        new_shape = [int(i) for i in (0, *(dset.shape[1:]))]
        self._resize_dataset(dset, new_shape)


    def _ensure_frame(self, dset, index):
        '''
        Ensures that `dset` has at least `index` points in its first dimension.

        Args:
            dset: Dataset object, dependent on the backend

            index: An integer counting the objects (starting with 1?)

        '''
        idiff = (index-dset.shape[0]) + 1

        if idiff == 1:
            shape_diff = 1
            
        elif idiff > 1:
            if '+' in self._frame_mode:
                shape_diff = idiff
            else:
                raise RefuseSparse(f'msg="Sparse writing requires top-up flag" '
                                   f'index={index} frames={dset.shape[0]} ')
            
        elif idiff <= 0:
            if 'a' in self._frame_mode: # append mode
                raise RefuseOverwrite(f'msg="Overwrite not permitted" index={index} '
                                      f'frames={dset.shape[0]} dataset="{dset.name}" '
                                      f'file="{self._recurse_obj_fname(dset)}"')

            elif 'o' in self._frame_mode or 'O' in self._frame_mode: # overwrite mode
                shape_diff = 0

            elif 'x' in self._frame_mode: # reset mode
                logger.info(f'msg="Resetting dataset" index={index} '
                            f'frames={dset.shape[0]} dataset="{dset.name}"')
                self._reset_dataset(dset)

                shape_diff = index+1

            else:
                raise RuntimeError(f'BUG: No shape diff for mode: {self._frame_mode}')

        if shape_diff > 0:
            new_shape = [int(i) for i in (dset.shape[0]+shape_diff, *(dset.shape[1:]))]
            logger.debug(f'msg="Resizing dataset" '
                         f'size={dset.shape[0]} '
                         f'index={index} diff={shape_diff} idiff={idiff} '
                         f'shape={new_shape}')
            self._resize_dataset(dset, new_shape)


    async def push_frame(self, index, **data):
        '''
        Received for every data point in a scan.

        Args:
            index: Integer to determine the position of the current
              data point in the overall (larger) dataset. This should
              always point to the next available index in the first
              dimention, i.e. if the existing dataset has dimension
              (37, N, M), with the current data point to become
              the 38th, then `index` needs to be 37 (counting starts
              at 0).

              Behavior if the index doesn't match the shape is defined
              by the class-global parameter `mode`.

              If this is `None`, the next free slot (i.e. current_size+1)
              is used, essentially making the operation always succeed
              as if "append" mode were active.

              If this is a negative number (e.g. -1) and the dataset
              doesn't exist, it is created according to shape
              specification of the input data, but with length 0
              (i.e. no points are copied).
              
            **data: dictionary from data name ("tag") to data arrays
        '''

        for key,container in data.items():
            dset_path = self._stor_dspath[key]
            grp, dset_only = self._split_dset_path(dset_path)
            self.ensure_group(grp)
            dset = self.ensure_dataset(dset_path,
                                       container.dtype,
                                       container.shape)

            if index is None:
                index = container.shape[0]

            if index >= 0:
                self._ensure_frame(dset, index)
                item = np_array(container)
                dset[index] = item


    def __getitem__(self, tag):
        '''
        Returns a handle to the object internally known as `tag`, according
        to the path spec specified at `__init__` time. This is generally
        supposed to be a dataset.
        '''
        dspath = self._stor_dspath[tag]
        obj = self.outer_object(tag)

        try:
            return obj[dspath]
        except KeyError:
            if not self.readonly:
                return self.ensure_dataset(dspath)
            raise
