"""
Preserve, for when pickling is just a bit too intense

Preserve is a simple framework to painlessly convert your python objects and
data structures into a nested structure containing only primitive types.

While intended for use in serialisation, Preserve stops short of the actual
serialisation bit - you can then use whatever method you prefer to dump
the nested primitive structure to JSON, TOML, or some other message packing
or storage system - most of which handle all primitive types by default.

"""

from enum import Enum, auto
import inspect


class RestoreMethod(Enum):
    DIRECT = auto()
    INIT = auto()
    CLASS_METHOD = auto()


class_key = "<_jam>"
"""
Contains the reserved dict key used to indicate that the dict it is in
should be restored to a class instance, not just treated as a dict.

Set to ``<_jam>`` by default. If changed externally, must be set before
any ``@preservable`` decorators
"""


preservables = {}
"""
A dict of classes marked as ``@preservable``, used to restore them back
to class instances.
"""


def preservable(cls, restore_method=RestoreMethod.DIRECT, name=None):
    """
    Decorator to mark class as preservable and keep track of associated names
    and classes.

    Args:
        restore_method: One of the available preserve.RestoreMethod values.
            Sets the method used for restoring this class. Defaults to
            ``DIRECT``, skipping the ``__init__`` method and setting all
            attributes as they were.
        name: The string used to indentify this class in the preserved dict.
            Must be unique among all ``@preservable`` classes. Defaults to the
            class name if left as None.
    """
    if name is None:
        cls._preserve_name = cls.__name__
    else:
        if isinstance(name, str):
            raise Exception("Preservable name must be a string")
        cls._preserve_name = name
    cls._restore_method = restore_method

    if cls._preserve_name in preservables:
        raise Exception("Duplicate preservable class name "+cls._preserve_name)
    preservables[cls._preserve_name] = cls

    def _preserve(self):
        dict_jam = _preserve_dict(vars(self))
        dict_jam[class_key] = self._preserve_name
        return dict_jam

    cls.preserve = _preserve
    return cls


def preserve(target_obj):
    """
    Preserve ``target_obj``, running through its contents recursively.

    Args:
        target_obj: The object to be preserved. This object and all its nested
            contents must either be primitive types or objects of a
            ``@preservable`` class.

    Returns:
        The preserved data structure - a nested structure containing only primitive
            types.
    """
    # If it's a primitive, store it. If it's a dict or list, recursively preserve that.
    # If it's an instance of another preservable class, call its .preserve() method.
    if isinstance(target_obj, (str, int, float, bool, type(None))):
        return target_obj
    elif isinstance(target_obj, dict):
        return _preserve_dict(target_obj)
    elif isinstance(target_obj, list):
        list_jam = []
        for val in target_obj:
            list_jam.append(preserve(val))
        return list_jam
    elif hasattr(target_obj, "_preserve_name"):
        return target_obj.preserve()
    else:
        raise Exception("Object "+str(target_obj)+" is not preservable")


def _preserve_dict(target_dict):
    dict_jam = {}
    for k, val in target_dict.items():
        if not isinstance(k, str):
            raise Exception("Non-string dictionary keys are not preservable")
        if k == class_key:
            raise Exception("Key "+class_key+" is reserved for internal use")
        dict_jam[k] = preserve(val)
    return dict_jam


def restore(obj_jam):
    """
    Restore the result of ``preserve()`` back into its original form. Will
    recursively scan the data structure and restore any
    ``@preservable`` classes according to their ``restore_method``.

    Args:
        obj_jam: The data structure to restore, usually the result of a
            ``preserve()`` call.
    """
    if isinstance(obj_jam, (str, int, float, bool, type(None))):
        return obj_jam
    elif isinstance(obj_jam, dict):
        return _restore_dict(obj_jam)
    elif isinstance(obj_jam, list):
        restored_list = []
        for val in obj_jam:
            restored_list.append(restore(val))
        return restored_list
    else:
        raise Exception("Object "+str(obj_jam)+" is not restorable")


def _restore_dict(dict_jam):
    restored_dict = {}
    for k, val in dict_jam.items():
        if not isinstance(k, str):
            raise Exception("Non-string dictionary keys are not restorable")
        if k != class_key:
            restored_dict[k] = restore(val)

    # Check if this is an object that needs to be restored back to a class instance
    if class_key in dict_jam:
        if dict_jam[class_key] not in preservables:
            raise Exception("Class "+dict_jam[class_key]+" has not been decorated as preservable")
        f_class = preservables[dict_jam[class_key]]
        # If DIRECT, skip __init__ and set attributes back directly
        if f_class._restore_method == RestoreMethod.DIRECT:
            restored_instance = f_class.__new__(f_class)
            restored_instance.__dict__.update(restored_dict)
        # if INIT, pass all attributes as keywords to __init__ method
        elif f_class._restore_method == RestoreMethod.INIT:
            restored_instance = f_class(**restored_dict)
        # IF CLASS_METHOD, pass all attributes as keyword aguments to classmethod "unpack()"
        elif f_class._restore_method == RestoreMethod.CLASS_METHOD:
            if inspect.ismethod(getattr(f_class, "restore", None)):
                restored_instance = f_class.restore(**restored_dict)
            else:
                raise Exception("Class "+str(f_class)+" does not have classmethod 'restore()'")
        else:
            raise Exception("Class _restore_method " +
                            str(f_class._restore_method)+" is not supported")

        return restored_instance
    else:
        return restored_dict