# encoding: utf-8
"""
# load_python.py

Handlers for dumping and loading built-in python types.
NB: As these are for built-in types, they are critical to the functioning of hickle.

"""

import six
from hickle.helpers import get_type_and_data

try:
    from exceptions import Exception
except ImportError:
    pass        # above imports will fail in python3

try:
    ModuleNotFoundError  # This fails on Py3.5 and below
except NameError:
    ModuleNotFoundError = ImportError

import h5py as h5


def get_py3_string_type(h_node):
    """ Helper function to return the python string type for items in a list.

    Notes:
        Py3 string handling is a bit funky and doesn't play too nicely with HDF5.
        We needed to add metadata to say if the strings in a list started off as
        bytes, string, etc. This helper loads

    """
    try:
        py_type = h_node.attrs["py3_string_type"][0]
        return py_type
    except:
        return None

def create_listlike_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Dumper for list, set, tuple

    Args:
        py_obj: python object to dump; should be list-like
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    dtype = str(type(py_obj))
    obj = list(py_obj)

    # h5py does not handle Py3 'str' objects well. Need to catch this
    # Only need to check first element as this method
    # is only called if all elements have same dtype
    py3_str_type = None
    if type(obj[0]) in (str, bytes):
        py3_str_type = bytes(str(type(obj[0])), 'ascii')

    if type(obj[0]) is str:
        #print(py3_str_type)
        #print(obj, "HERE")
        obj = [bytes(oo, 'utf8') for oo in obj]
        #print(obj, "HERE")


    d = h_group.create_dataset('data_%i' % call_id, data=obj, **kwargs)
    d.attrs["type"] = [bytes(dtype, 'ascii')]

    # Need to add some metadata to aid in unpickling if it's a string type
    if py3_str_type is not None:
        d.attrs["py3_string_type"] = [py3_str_type]



def create_python_dtype_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ dumps a python dtype object to h5py file

    Args:
        py_obj: python object to dump; should be a python type (int, float, bool etc)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    # kwarg compression etc does not work on scalars
    d = h_group.create_dataset('data_%i' % call_id, data=py_obj,
                               dtype=type(py_obj))     #, **kwargs)
    d.attrs["type"] = [b'python_dtype']
    d.attrs['python_subdtype'] = bytes(str(type(py_obj)), 'ascii')


def create_stringlike_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ dumps a list object to h5py file

    Args:
        py_obj: python object to dump; should be string-like (unicode or string)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    if isinstance(py_obj, bytes):
        d = h_group.create_dataset('data_%i' % call_id, data=[py_obj], **kwargs)
        d.attrs["type"] = [b'bytes']
    elif isinstance(py_obj, str):
        dt = h5.special_dtype(vlen=str)
        dset = h_group.create_dataset('data_%i' % call_id, shape=(1, ), dtype=dt, **kwargs)
        dset[0] = py_obj
        dset.attrs['type'] = [b'string']

def create_none_dataset(py_obj, h_group, call_id=0, **kwargs):
    """ Dump None type to file

    Args:
        py_obj: python object to dump; must be None object
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    d = h_group.create_dataset('data_%i' % call_id, data=[0], **kwargs)
    d.attrs["type"] = [b'none']


def load_list_dataset(h_node):
    py_type, data = get_type_and_data(h_node)
    py3_str_type = get_py3_string_type(h_node)

    if py3_str_type == b"<class 'bytes'>":
        # Yuck. Convert numpy._bytes -> str -> bytes
        return [bytes(str(item, 'utf8'), 'utf8') for item in data]
    if py3_str_type == b"<class 'str'>":
        return [str(item, 'utf8') for item in data]
    else:
        return list(data)

def load_tuple_dataset(h_node):
    data = load_list_dataset(h_node)
    return tuple(data)

def load_set_dataset(h_node):
    data = load_list_dataset(h_node)
    return set(data)

def load_bytes_dataset(h_node):
    py_type, data = get_type_and_data(h_node)
    return bytes(data[0])

def load_string_dataset(h_node):
    py_type, data = get_type_and_data(h_node)
    return str(data[0])

def load_unicode_dataset(h_node):
    py_type, data = get_type_and_data(h_node)
    return unicode(data[0])

def load_none_dataset(h_node):
    return None

def load_pickled_data(h_node):
    py_type, data = get_type_and_data(h_node)
    try:
        import cPickle as pickle
    except ModuleNotFoundError:
        import pickle
    return pickle.loads(data[0])


def load_python_dtype_dataset(h_node):
    py_type, data = get_type_and_data(h_node)
    subtype = h_node.attrs["python_subdtype"]
    type_dict = {
        b"<class 'int'>": int,
        b"<class 'float'>": float,
        b"<class 'bool'>": bool,
        b"<class 'complex'>": complex
    }

    tcast = type_dict.get(subtype)
    return tcast(data)



types_dict = {
    list:        create_listlike_dataset,
    tuple:       create_listlike_dataset,
    set:         create_listlike_dataset,
    bytes:         create_stringlike_dataset,
    str:           create_stringlike_dataset,
    #bytearray:     create_stringlike_dataset,
    int:         create_python_dtype_dataset,
    float:       create_python_dtype_dataset,
    bool:        create_python_dtype_dataset,
    complex:     create_python_dtype_dataset,
    type(None):    create_none_dataset,
}

hkl_types_dict = {
    b"<class 'list'>"  : load_list_dataset,
    b"<class 'tuple'>" : load_tuple_dataset,
    b"<class 'set'>"   : load_set_dataset,
    b"bytes"           : load_bytes_dataset,
    b"python_dtype"   : load_python_dtype_dataset,
    b"string"         : load_string_dataset,
    b"pickle"         : load_pickled_data,
    b"none"           : load_none_dataset,
}