# data pool collects data from data handling script
# provides data to experiment script and display

import sys
import tables
import collections
import threading
import traceback
import io
import numpy

class DataPool(collections.abc.MutableMapping):
    """
    dictionary with sending change events
    """

    # supports tranlation from dictionary keys to pytables hdf node names
    # taken from: Python Ref Manual Section 2.3: Identifiers and keywords
    # things are always prefixed by "dir_" or "dict_"
    translation_table=""
    for i in range(256):
        c=chr(i)
        if (c>="a" and c<="z") or \
           (c>="A" and c<="Z") or \
           (c>="0" and c<="9"):
            translation_table+=c
        else:
            translation_table+="_"

    class Event:
        access=0
        updated_value=1
        new_key=2
        deleted_key=3
        destroy=4

        def __init__(self, what, subject="", origin=None):
            self.what=what
            self.subject=subject
            self.origin=origin

        def __repr__(self):
            return "<DataPool.Event origin=%s what=%d subject='%s'>"%(self.origin, self.what,self.subject)

        def copy(self):
            return DataPool.Event(self.what+0, self.subject+"", self.origin)

    def __init__(self):
        self.__mydict={}
        self.__dictlock=threading.Lock()
        self.__registered_listeners=[]
        self.__dirty_keys=set()

    def __getitem__(self, name):
        try:
            self.__dictlock.acquire()
            return self.__mydict[name]
        finally:
            self.__dictlock.release()

    def __setitem__(self, name, value):
        try:
            self.__dictlock.acquire()
            if name in self.__mydict:
                e=DataPool.Event(DataPool.Event.updated_value,name,self)
            else:
                e=DataPool.Event(DataPool.Event.new_key, name,self)
            self.__mydict[name]=value
            self.__dirty_keys.add(name)
        finally:
            self.__dictlock.release()
        self.__send_event(e)


    def __delitem__(self, name):
        try:
            self.__dictlock.acquire()
            del self.__mydict[name]
            self.__dirty_keys.add(name)
        finally:
            self.__dictlock.release()
        self.__send_event(DataPool.Event(DataPool.Event.deleted_key,name,self))

    def __iter__(self):
        try:
            self.__dictlock.acquire()
            return iter(self.__mydict)
        finally:
            self.__dictlock.release()

    def __len__(self):
        try:
            self.__dictlock.acquire()
            return len(self.__mydict)
        finally:
            self.__dictlock.release()

    def keys(self):
        try:
            self.__dictlock.acquire()
            return list(self.__mydict.keys())
        finally:
            self.__dictlock.release()

    def __send_event(self, _event):
        for listeners in self.__registered_listeners:
            listeners(_event.copy())

    def __del__(self):
        self.__send_event(DataPool.Event(DataPool.Event.destroy))
        self.__registered_listeners=None

    def write_hdf5(self,hdffile,where="/",name="data_pool", complib=None, complevel=None, incremental=False):

        """
        Writes the content of a DataPool object to an HDF5 file using the `tables` library.

        This method serializes the entire content of the data pool into an HDF5 file, creating
        appropriate groups and arrays based on the keys and values stored in the data pool.
        Nested storage paths are supported, and group or array names are
        converted into valid HDF5 names. Paths are prepended with "dir_", results with "dict_".
        If duplicate names arise, extensions (like `_001`, `_002`)
        are appended to ensure uniqueness.

        Parameters:
        hdffile : str | bytes | tables.File
            Path to the HDF5 file or an already open `tables.File` object where the
            data should be written. Strings or bytes are interpreted as file paths.

        where : str, optional
            The root HDF5 group path under which the data pool will be stored.
            Defaults to "/".

        name : str, optional
            The name of the root HDF5 group under which the data pool is saved.
            Defaults to "data_pool".

        complib : str, optional
            Compression library to be used when creating HDF5 arrays. Optional, defaults
            to None, meaning no compression is applied.

        complevel : int, optional
            Compression level to use when applying compression. Has no effect if `complib`
            is set to None.

        Raises:
        Exception
            If `hdffile` is neither a string nor a valid `tables.File` object.

        Notes:
        - This method assumes that the content of the data pool does not change during the
          write operation.
        - Keys in the data pool starting with "__" (double underscores) are skipped.
        - If a key contains nested paths (e.g., "group1/subgroup2/item"), intermediate
          groups are created as necessary.
        - Group or array names are converted to valid HDF5 names, and name conflicts
          are resolved by appending extensions (e.g., "_001").
        - Errors during the writing process are logged to the console, and traceback
          information is displayed for debugging purposes.
        """
        if isinstance(hdffile, (bytes, str)):
            dump_file=tables.open_file(hdffile, mode="a")
        elif isinstance(hdffile,tables.File):
            dump_file=hdffile
        else:
            raise Exception("expecting hdffile or string")

        if where in dump_file and name in dump_file.get_node(where):
            dump_group = dump_file.get_node(where, name)
        else:
            dump_group = dump_file.create_group(where, name, "DAMARIS data pool")

        self.__dictlock.acquire()
        if incremental:
            dict_keys=list(self.__dirty_keys)
        else:
            # Full sync: current keys + any other dirty keys (e.g. deletions)
            dict_keys=list(set(self.__mydict.keys()) | self.__dirty_keys)
        self.__dictlock.release()

        try:
            for key in dict_keys:
                if key.startswith("__"):
                    continue
                dump_dir=dump_group
                # walk along the given path and create groups if necessary
                namelist = key.split("/")
                for part in namelist[:-1]:
                    dir_part="dir_"+str(part).translate(DataPool.translation_table)
                    if dir_part not in dump_dir:
                        dump_dir=dump_file.create_group(dump_dir,name=dir_part,title=part)
                    else:
                        if dump_dir._v_children[dir_part]._v_title==part:
                            dump_dir=dump_dir._v_children[dir_part]
                        else:
                            extension_count=0
                            while dir_part+"_%03d"%extension_count in dump_dir:
                                extension_count+=1
                            dump_dir=dump_file.create_group(dump_dir,
                                                           name=dir_part+"_%03d"%extension_count,
                                                           title=part)

                # convert last part of key to a valid name
                group_keyname="dict_"+str(namelist[-1]).translate(DataPool.translation_table)

                # find existing node for this key
                if group_keyname in dump_dir:
                    if dump_dir._v_children[group_keyname]._v_title == key:
                        dump_file.remove_node(dump_dir, group_keyname, recursive=True)
                    else:
                        extension_count = 0
                        while group_keyname + "_%03d" % extension_count in dump_dir:
                            suffix_name = group_keyname + "_%03d" % extension_count
                            if dump_dir._v_children[suffix_name]._v_title == key:
                                dump_file.remove_node(dump_dir, suffix_name, recursive=True)
                                group_keyname = suffix_name
                                break
                            extension_count += 1
                        else:
                            # truly new key with collision
                            group_keyname += "_%03d" % extension_count

                self.__dictlock.acquire()
                if key not in self.__mydict:
                    # Key was deleted
                    if key in self.__dirty_keys:
                        self.__dirty_keys.remove(key)
                    self.__dictlock.release()
                    continue
                value=self.__mydict[key]
                if key in self.__dirty_keys:
                    self.__dirty_keys.remove(key)
                self.__dictlock.release()
                # now write data, assuming, the object is constant during write operation
                if hasattr(value, "write_to_hdf"):
                    try:
                        value.write_to_hdf(hdffile=dump_file,
                                           where=dump_dir,
                                           name=group_keyname,
                                           title=key,
                                           complib=complib,
                                           complevel=complevel)
                    except Exception as e:
                        print("failed to write data_pool[\"%s\"]: %s"%(key,str(e)))
                        traceback_file=io.StringIO()
                        traceback.print_tb(sys.exc_info()[2], None, traceback_file)
                        print("detailed traceback: %s\n"%str(e)+traceback_file.getvalue())
                        traceback_file=None
                elif isinstance(value, (numpy.ndarray, int, float, str, bytes)):
                    try:
                        obj_to_save = value
                        if isinstance(value, str):
                            obj_to_save = value.encode('utf-8')
                        dump_file.create_array(dump_dir,
                                              name=group_keyname,
                                              obj=obj_to_save,
                                              title=key)
                    except Exception as e:
                        print("failed to write data_pool[\"%s\"]: %s"%(key,str(e)))
                else:
                    print("don't know how to store data_pool[\"%s\"] (type: %s)"%(key, type(value)))
                value=None

        finally:
            dump_group=None
            if isinstance(hdffile, (bytes, str)):
                dump_file.close()
            dump_file=None

    def read_from_hdf(self, hdffile, where="/data_pool"):
        """
        Read specified data from HDF5 file and populate the DataPool.
        
        Parameters:
        - hdffile: HDF5 file object or filename (bytes/str)
        - where: Location within HDF5 file to read from (default: "/data_pool")
        """
        # Handle both filename and file object
        if type(hdffile) is bytes or isinstance(hdffile, str):
            hdf_file = tables.open_file(hdffile, mode="r")
            close_file = True
        elif isinstance(hdffile, tables.File):
            hdf_file = hdffile
            close_file = False
        else:
            raise Exception("expecting hdffile or string")
        
        try:
            # Navigate to the data pool group
            if where not in hdf_file:
                raise Exception(f"HDF5 location '{where}' not found")
            
            data_pool_group = hdf_file.get_node(where)
            
            # Recursively process the group structure
            self._read_hdf_group(data_pool_group, "")
            
            # Clear dirty keys as they are now in sync with HDF5
            self.__dictlock.acquire()
            self.__dirty_keys.clear()
            self.__dictlock.release()
            
        finally:
            if close_file:
                hdf_file.close()

    def _read_hdf_group(self, group, path_prefix):
        """
        Recursively read HDF5 group and populate DataPool.
        
        Parameters:
        - group: HDF5 group object
        - path_prefix: Current path in DataPool hierarchy
        """
        # Process all nodes in this group
        for node in group:
            # Get the node name (last part of the path)
            node_name = node._v_name
            
            # Try to get the original name from title
            title = node._v_title
            
            if title:
                # For DAMARIS objects, the title is usually the full path
                # For intermediate groups, it's just the part name
                if "/" in title:
                    full_path = title
                else:
                    if path_prefix:
                        full_path = f"{path_prefix}/{title}"
                    else:
                        full_path = title
            else:
                # Strip dir_ and dict_ prefixes from node name
                clean_name = self._clean_hdf_name(node_name)
                
                # Build the full path in DataPool hierarchy
                if path_prefix:
                    full_path = f"{path_prefix}/{clean_name}"
                else:
                    full_path = clean_name
            
            full_path = full_path.replace("//", "/")
            
            if isinstance(node, (tables.Group, tables.Table)):
                # Check if this is a known DAMARIS object type
                if hasattr(node._v_attrs, 'damaris_type'):
                    damaris_type = node._v_attrs.damaris_type
                    
                    # Handle different DAMARIS object types
                    if damaris_type == "Accumulation":
                        from damaris.data.Accumulation import read_from_hdf
                        obj = read_from_hdf(node)
                        if obj is not None:
                            self[full_path] = obj
                    
                    elif damaris_type == "ADC_Result":
                        from damaris.data.ADC_Result import read_from_hdf
                        obj = read_from_hdf(node)
                        if obj is not None:
                            self[full_path] = obj
                            
                    elif damaris_type == "MeasurementResult":
                        from damaris.data.MeasurementResult import read_from_hdf
                        obj = read_from_hdf(node)
                        if obj is not None:
                            self[full_path] = obj

                    elif damaris_type == "TemperatureResult":
                        from damaris.data.Temperature import read_from_hdf
                        obj = read_from_hdf(node)
                        if obj is not None:
                            self[full_path] = obj

                    # Skip further processing of this group since we handled it
                    continue
                
                # Recursively process sub-groups
                if isinstance(node, tables.Group):
                    self._read_hdf_group(node, full_path)
                
            elif isinstance(node, (tables.Array, tables.CArray, tables.EArray, tables.Table)):
                # Handle simple array data
                try:
                    data = node.read()
                    # If it's a scalar array of bytes, decode it to string
                    if hasattr(data, 'decode'):
                        try:
                            data = data.decode('utf-8')
                        except:
                            pass
                    elif isinstance(data, numpy.ndarray) and data.dtype.kind in ('S', 'V'):
                        # For numpy byte arrays, try to decode if they are scalar or small
                        if data.size == 1:
                            try:
                                decoded = data.item().decode('utf-8')
                                data = decoded
                            except:
                                pass
                    self[full_path] = data
                except Exception as e:
                    print(f"Warning: Could not read array {full_path}: {e}")
            
            # Note: Other node types (Table, etc.) would be handled here as needed
    
    def _clean_hdf_name(self, name):
        """
        Clean HDF5 node name by removing dir_ and dict_ prefixes.
        
        Parameters:
        - name: HDF5 node name (may have dir_ or dict_ prefix)
        
        Returns:
        - Cleaned name with prefixes removed
        """
        if name.startswith("dir_"):
            return name[4:]  # Remove "dir_" prefix
        elif name.startswith("dict_"):
            return name[5:]  # Remove "dict_" prefix
        else:
            return name

    @classmethod
    def load_hdf5(cls, filename):
        """
        Load a complete DAMARIS HDF5 file and return DataPool and metadata, including scripts.
        
        Parameters:
        - filename: Path to HDF5 file
        
        Returns:
        - tuple: (DataPool instance, metadata dict)
        
        metadata dict contains:
        - 'scripts': dict with experiment scripts and spool info
        - 'log': log data if available
        - 'timeline': timeline data if available
        """
        # Validate file
        if not tables.is_pytables_file(filename):
            raise Exception(f"File {filename} is not a valid PyTables HDF5 file")
        
        # Create DataPool and load data
        data_pool = cls()
        
        try:
            hdf_file = tables.open_file(filename, mode="r")
            
            # Load main data pool
            if "/data_pool" in hdf_file:
                data_pool.read_from_hdf(hdf_file, where="/data_pool")
            
            # Extract metadata
            metadata = {}
            
            # Load scripts
            if "/scripts" in hdf_file:
                scripts_group = hdf_file.get_node("/scripts")
                scripts_data = {}
                # Use _v_children to get node names safely
                if hasattr(scripts_group, '_v_children'):
                    for node_name, node in scripts_group._v_children.items():
                        try:
                            if hasattr(node, 'read'):
                                data = node.read()
                                # Handle numpy string arrays and other types
                                if hasattr(data, 'decode'):
                                    # bytes or numpy bytes
                                    scripts_data[node_name] = data.decode('utf-8', errors='replace')
                                elif hasattr(data, 'tolist'):
                                    # numpy array
                                    scripts_data[node_name] = str(data.tolist())
                                elif isinstance(data, (str, bytes)):
                                    # string or bytes
                                    if isinstance(data, bytes):
                                        scripts_data[node_name] = data.decode('utf-8', errors='replace')
                                    else:
                                        scripts_data[node_name] = data
                                else:
                                    scripts_data[node_name] = str(data)
                        except Exception as e:
                            print(f"Warning: Could not read script {node_name}: {e}")
                metadata['scripts'] = scripts_data
            
            # Load log
            if "/log" in hdf_file:
                try:
                    log_node = hdf_file.get_node("/log")
                    metadata['log'] = log_node.read()
                except Exception as e:
                    print(f"Warning: Could not read log: {e}")
            
            # Load timeline
            if "/timeline" in hdf_file:
                try:
                    timeline_node = hdf_file.get_node("/timeline")
                    metadata['timeline'] = timeline_node.read()
                except Exception as e:
                    print(f"Warning: Could not read timeline: {e}")
            
            return data_pool, metadata
            
        except Exception as e:
            raise Exception(f"Failed to load HDF5 file {filename}: {e}")
        finally:
            if 'hdf_file' in locals():
                hdf_file.close()

    def register_listener(self, listening_function):
        self.__registered_listeners.append(listening_function)

    def unregister_listener(self, listening_function):
        if listening_function in self.__registered_listeners:
            self.__registered_listeners.remove(listening_function)
