"""write_hdf5.py: Class to write a resqml hdf5 file."""

version = '26th August 2021'

# Nexus is a registered trademark of the Halliburton Company

# write_hdf5.py
# approach is to register the datasets (arrays) to be written; then write everything in a separate, single function call

import logging

log = logging.getLogger(__name__)
log.debug('write_hdf5.py version ' + version)

import os
import numpy as np

import h5py
import resqpy.olio.uuid as bu

resqml_path_head = '/RESQML/'  # note: latest fesapi code uses RESQML20
write_bool_as_int8 = True  # Nexus read fails if bool used as hdf5 element dtype; also better for NullValue handling
write_int_as_int32 = True  # only applies if registered dtype is None


class H5Register():
   """Class for registering arrays and then writing to an hdf5 file."""

   def __init__(self, model):
      """Create a new, empty register of arrays to be written to an hdf5 file."""

      self.dataset_dict = {}  # dictionary mapping from (object_uuid, group_tail) to (numpy_array, dtype)
      self.hdf5_path_dict = {}  # dictionary optionally mapping from (object_uuid, group_tail) to hdf5 internal path
      self.model = model

   def register_dataset(self, object_uuid, group_tail, a, dtype = None, hdf5_path = None):
      """Register an array to be included as a dataset in the hdf5 file.

      arguments:
         object_uuid (uuid.UUID): the uuid of the object (part) that this array is for
         group_tail (string): the remainder of the hdf5 internal path (following RESQML and
            uuid elements)
         a (numpy array): the dataset (array) to be registered for writing
         dtype (type or string): the type of the individual elements within the dataset
         hdf5_path (string, optional): if present, a full hdf5 internal path to use instead of
            the default generated from the uuid

      returns:
         None

      note:
         several arrays might belong to the same object
      """

      #     print('registering dataset with uuid ' + str(object_uuid) + ' and group tail ' + group_tail)
      assert (len(group_tail) > 0)
      assert a is not None
      assert isinstance(a, np.ndarray)
      if group_tail[0] == '/':
         group_tail = group_tail[1:]
      if group_tail[-1] == '/':
         group_tail = group_tail[:-1]
      if (object_uuid, group_tail) in self.dataset_dict.keys():
         pass  # todo: warn of re-registration?
      self.dataset_dict[(object_uuid, group_tail)] = (a, dtype)
      if hdf5_path:
         self.hdf5_path_dict[(object_uuid, group_tail)] = hdf5_path

   def write_fp(self, fp):
      """Write or append to an hdf5 file, writing the pre-registered datasets (arrays).

      arguments:
         fp: an already open h5py._hl.files.File object

      returns:
         None

      note:
         the file handle fp must have been opened with mode 'w' or 'a'
      """

      # note: in resqml, an established hdf5 file has a uuid and should therefore be immutable
      #       this function allows appending to any hdf5 file; calling code should set a new uuid when needed
      assert (fp is not None)
      for (object_uuid, group_tail) in self.dataset_dict.keys():
         if (object_uuid, group_tail) in self.hdf5_path_dict.keys():
            hdf5_path = self.hdf5_path_dict[(object_uuid, group_tail)]
         else:
            hdf5_path = resqml_path_head + str(object_uuid) + '/' + group_tail
         (a, dtype) = self.dataset_dict[(object_uuid, group_tail)]
         if dtype is None:
            dtype = a.dtype
            if write_int_as_int32 and str(dtype) == 'int64':
               dtype = 'int32'
         if write_bool_as_int8 and str(dtype).lower().startswith('bool'):
            dtype = 'int8'
         log.debug('Writing hdf5 dataset ' + hdf5_path + ' of size ' + str(a.size) + ' type ' + str(dtype))
         fp.create_dataset(hdf5_path, data = a, dtype = dtype)

   def write(self, file = None, mode = 'w', release_after = True):
      """Create or append to an hdf5 file, writing the pre-registered datasets (arrays).

      arguments:
         file: either a string being the file path, or an already open h5py._hl.files.File object;
            if None (recommended), the file is opened through the model object's hdf5 management
            functions
         mode (string, default 'w'): the mode to open the file in; only relevant if file is a path;
            must be 'w' or 'a' for (over)write or append

      returns:
         None
      """

      # note: in resqml, an established hdf5 file has a uuid and should therefore be immutable
      #       this function allows appending to any hdf5 file;
      #       strictly, calling code should set a new uuid when needed, in practice not essential
      if len(self.dataset_dict) == 0:
         return
      if file is None:
         file = self.model.h5_access(mode = mode)
      elif isinstance(file, str):
         file = self.model.h5_access(mode = mode, file_path = file)
      if mode == 'a' and isinstance(file, str) and not os.path.exists(file):
         mode = 'w'
      assert isinstance(file, h5py._hl.files.File)
      self.write_fp(file)
      if release_after:
         self.model.h5_release()


def copy_h5(file_in, file_out, uuid_inclusion_list = None, uuid_exclusion_list = None, mode = 'w'):
   """Create a copy of an hdf5, optionally including or excluding arrays with specified uuids.

   arguments:
      file_in (string): path of existing hdf5 file to be duplicated
      file_out (string): path of output hdf5 file to be created or appended to (see mode)
      uuid_inclusion_list (list of uuid.UUID, optional): if present, the uuids to be included
         in the output file
      uuid_exclusion_list (list of uuid.UUID, optional): if present, the uuids to be excluded
         from the output file
      mode (string, default 'w'): mode to open output file with; must be 'w' or 'a' for
         (over)write or append respectively

   returns:
      number of hdf5 groups (uuid's) copied

   notes:
      at most one of uuid_inclusion_list and uuid_exclusion_list should be passed;
      if neither are passed, all the datasets (arrays) in the input file are copied to the
      output file
   """

   #  note: if both inclusion and exclusion lists are present, exclusion list is ignored
   assert file_out != file_in, 'identical input and output files specified for hdf5 copy'
   assert uuid_inclusion_list is None or uuid_exclusion_list is None,  \
      'inclusion and exclusion lists both specified for hdf5 copy; at most one allowed'
   checking_uuid = uuid_inclusion_list is not None or uuid_exclusion_list is not None
   assert mode in ['w', 'a']
   copy_count = 0
   with h5py.File(file_out, mode) as fp_out:
      assert fp_out is not None, 'failed to open output hdf5 file: ' + file_out
      with h5py.File(file_in, 'r') as fp_in:
         assert fp_in is not None, 'failed to open input hdf5 file: ' + file_in
         main_group_in = fp_in['RESQML']
         assert main_group_in is not None, 'failed to find RESQML group in hdf5 file: ' + file_in
         if mode == 'w':
            main_group_out = fp_out.create_group('RESQML')
         elif mode == 'a':
            try:
               main_group_out = fp_out['RESQML']
            except Exception:
               main_group_out = fp_out.create_group('RESQML')
         else:
            main_group_out = fp_out['RESQML']
         for group in main_group_in:
            if checking_uuid:
               uuid = bu.uuid_from_string(group)
               if uuid_inclusion_list is not None:
                  if uuid not in uuid_inclusion_list:
                     if uuid is None:
                        log.warning('RESQML group name in hdf5 file does not start with a uuid, skipping: ' +
                                    str(group))
                     continue
               else:  # uuid_exclusion_list is not None
                  if uuid in uuid_exclusion_list:
                     continue
                  if uuid is None:  # will still be copied
                     log.warning('RESQML group name in hdf5 file does not start with a uuid: ' + str(group))
            if group in main_group_out:
               log.warning('not copying hdf5 data due to pre-existence for: ' + str(group))
               continue
            log.debug('copying hdf5 data for uuid: ' + group)
            main_group_in.copy(group, main_group_out, expand_soft = True, expand_external = True, expand_refs = True)
            copy_count += 1
   return copy_count


def copy_h5_path_list(file_in, file_out, hdf5_path_list, mode = 'w'):
   """Create a copy of some hdf5 datasets (or groups), identified as a list of hdf5 internal paths.

   arguments:
      file_in (string): path of existing hdf5 file to be copied from
      file_out (string): path of output hdf5 file to be created or appended to (see mode)
      hdf5_path_list (list of string): the hdf5 internal paths of the datasets (or groups) to be copied
      mode (string, default 'w'): mode to open output file with; must be 'w' or 'a' for
         (over)write or append respectively

   returns:
      number of hdf5 datasets (or groups) copied
   """

   #  note: if both inclusion and exclusion lists are present, exclusion list is ignored
   assert file_out != file_in, 'identical input and output files specified for hdf5 copy'
   assert hdf5_path_list is not None
   assert mode in ['w', 'a']
   copy_count = 0
   with h5py.File(file_out, mode) as fp_out:
      assert fp_out is not None, f'failed to open output hdf5 file: {file_out}'
      with h5py.File(file_in, 'r') as fp_in:
         assert fp_in is not None, f'failed to open input hdf5 file: {file_in}'
         for path in hdf5_path_list:
            if path in fp_out:
               log.warning(f'not copying hdf5 data due to pre-existence for: {path}')
               continue
            assert path in fp_in, f'internal path {path} not found in hdf5 file {file_in}'
            log.debug(f'copying hdf5 data for: {path}')
            build = ''
            group_list = list(path.split(sep = '/'))
            assert len(group_list) > 1, f'no hdf5 group(s) in internal path {path}'
            for w in group_list[:-1]:
               if w:
                  build += '/' + w
                  if build not in fp_out:
                     fp_out.create_group(build)
            build += '/' + group_list[-1]
            fp_out.create_dataset(build, data = fp_in[path])
            #            fp_in.copy(path, fp_out[path], expand_soft = True, expand_external = True, expand_refs = True)
            copy_count += 1
   return copy_count


def change_uuid(file, old_uuid, new_uuid):
   """Changes hdf5 internal path (group name) for part, switching from old to new uuid.

   notes:
      this is low level functionality not usually called directly;
      the function assumes that hdf5 internal path names conform to the format that resqpy uses
      when writing data, namely /RESQML/uuid/tail...
   """

   assert file, 'hdf5 file name missing'
   assert old_uuid is not None and new_uuid is not None, 'missing uuid'

   def change_uuid_fp(fp, old_uuid, new_uuid):
      main_group = fp[resqml_path_head.strip('/')]
      old_group = main_group[str(old_uuid)]
      main_group[str(new_uuid)] = old_group
      del main_group[str(old_uuid)]

   if isinstance(file, h5py._hl.files.File):
      change_uuid_fp(file, old_uuid, new_uuid)
   else:
      assert isinstance(file, str)
      with h5py.File(file, 'r+') as fp:
         change_uuid_fp(fp, old_uuid, new_uuid)
