from enum import Enum, auto
import inspect

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

#freezedry, for when pickling is just a bit too intense

# The class key is a reserved dict key used to flag that the dict should be unpacked back out to a class instance
class_key = "<freezedried>"
# The Pack module stores some state from init to keep a list of valid packable classes
freezedryables = {}

# Decorator to mark class as packable and keep track of associated names and classes. When packed, the 
# special key string "<packable>" indicates what class the current dict should be unpacked to

# name argument is the string that will identify this class in a packed dict
def freezedryable(cls, rehydrate_method=RehydrateMethod.DIRECT, name=None):
    if name is None:
        cls._freezedry_name = cls.__name__
    else:
        if isinstance(name, str):
            raise Exception("freezedryable name must be a string")
        cls._freezedry_name = name
    cls._rehydrate_method = rehydrate_method

    if cls._freezedry_name in freezedryables:
        raise Exception("Duplicate freezedryable class name "+cls._freezedry_name)
    freezedryables[cls._freezedry_name] = cls

    def _freezedry(self):
        dried_dict=_freezedry_dict(vars(self))
        dried_dict[class_key]=self._freezedry_name
        return dried_dict

    cls.freezedry=_freezedry
    #setattr(cls, "freezedry", freezedry)
    return cls


def freezedry(hydrated_obj):
    # If it's a primitive, store it. If it's a dict or list, recursively freezedry that.
    # If it's an instance of another freezedryable class, call its .freezedry() method.
    if isinstance(hydrated_obj, (str, int, float, bool, type(None))):
        return hydrated_obj
    elif isinstance(hydrated_obj, dict):
        return _freezedry_dict(hydrated_obj)
    elif isinstance(hydrated_obj, list):
        dried_list = []
        for val in hydrated_obj:
            dried_list.append(freezedry(val))
        return dried_list
    elif hasattr(hydrated_obj, "_freezedry_name"):
        return hydrated_obj.freezedry()
    else:
        raise Exception("Object "+str(hydrated_obj)+" is not freezedryable")

def _freezedry_dict(hydrated_dict):
    dried_dict = {}
    for k,val in hydrated_dict.items():
        if not isinstance(k,str):
            raise Exception("Non-string dictionary keys are not freezedryable")
        if k == class_key:
            raise Exception("Key "+class_key+" is reserved for internal freezedry use")
        dried_dict[k]=freezedry(val)
    return dried_dict

def rehydrate(dried_obj):
    if isinstance(dried_obj, (str, int, float, bool, type(None))):
        return dried_obj
    elif isinstance(dried_obj, dict):
        return _rehydrate_dict(dried_obj)
    elif isinstance(dried_obj, list):
        hydrated_list = []
        for val in dried_obj:
            hydrated_list.append(rehydrate(val))
        return hydrated_list
    else:
        raise Exception("Object "+str(dried_obj)+" is not rehydrateable")

def _rehydrate_dict(dried_dict):
    hydrated_dict = {}
    for k,val in dried_dict.items():
        if not isinstance(k,str):
            raise Exception("Non-string dictionary keys are not rehydrateable")
        if k != class_key:
            hydrated_dict[k]=rehydrate(val)

    # Check if this is an object that needs to be unpacked back to an instance
    if class_key in dried_dict:
        if dried_dict[class_key] not in freezedryables:
            raise Exception("Class "+dried_dict[class_key]+" has not been decorated as freezedryable")
        f_class=freezedryables[dried_dict[class_key]]
        # If DIRECT, skip __init__ and set attributes back directly
        if f_class._rehydrate_method == RehydrateMethod.DIRECT:
            hydrated_instance = f_class.__new__(f_class)
            hydrated_instance.__dict__.update(hydrated_dict)
        #if INIT, pass all attributes as keywords to __init__ method
        elif f_class._rehydrate_method == RehydrateMethod.INIT:
            hydrated_instance = f_class(**hydrated_dict)
        # IF CLASS_METHOD, pass all attributes as keyword aguments to classmethod "unpack()"
        elif f_class._rehydrate_method == RehydrateMethod.CLASS_METHOD:
            if inspect.ismethod(getattr(f_class, "rehydrate", None)):
                hydrated_instance = f_class.rehydrate(**hydrated_dict)
            else:
                raise Exception("Class "+str(f_class)+" does not have classmethod 'rehydrate()'")
        else:
            raise Exception("Class _rehydrate_method "+str(f_class._rehydrate_method)+" is not supported")

        return hydrated_instance
    else:
        return hydrated_dict