"""
Here are all subprocess, threading etc related utilities,
most of them quite low level.
"""

from __future__ import annotations
from threading import Lock
import sys

import os
import io
import pickle
import types
import struct
import marshal
from importlib import import_module
import numpy

try:
    from _multiprocessing import Connection
except ImportError:
    from multiprocessing.connection import Connection

PY3 = sys.version_info[0] >= 3

if PY3:
    from io import BytesIO
else:
    # noinspection PyUnresolvedReferences,PyCompatibility
    from StringIO import StringIO as BytesIO


_abs_mod_file = os.path.abspath(__file__)


class SharedMem:
    class ShmException(Exception):
        pass

    class CCallException(ShmException):
        pass

    if sys.platform != "win32":
        import ctypes
        import ctypes.util

        libc_so = ctypes.util.find_library("c")
        libc = ctypes.CDLL(libc_so, use_errno=True)
        shm_key_t = ctypes.c_int
        IPC_PRIVATE = 0
        IPC_RMID = 0

        # int shmget(key_t key, size_t size, int shmflg);
        shmget = libc.shmget
        shmget.restype = ctypes.c_int
        shmget.argtypes = (shm_key_t, ctypes.c_size_t, ctypes.c_int)
        # void* shmat(int shmid, const void *shmaddr, int shmflg);
        shmat = libc.shmat
        shmat.restype = ctypes.c_void_p
        shmat.argtypes = (ctypes.c_int, ctypes.c_void_p, ctypes.c_int)
        # int shmdt(const void *shmaddr);
        shmdt = libc.shmdt
        shmdt.restype = ctypes.c_int
        shmdt.argtypes = (ctypes.c_void_p,)
        # int shmctl(int shmid, int cmd, struct shmid_ds *buf);
        shmctl = libc.shmctl
        shmctl.restype = ctypes.c_int
        shmctl.argtypes = (ctypes.c_int, ctypes.c_int, ctypes.c_void_p)
        # void* memcpy( void *dest, const void *src, size_t count );
        memcpy = libc.memcpy
        memcpy.restype = ctypes.c_void_p
        memcpy.argtypes = (ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t)

        @classmethod
        def check_ccall_error(cls, check, f):
            import ctypes

            if not check:
                errno = ctypes.get_errno()
                errstr = os.strerror(errno)
                raise cls.CCallException("SharedMem: %s failed with error %i (%s)" % (f, errno, errstr))

        @classmethod
        def is_shmget_functioning(cls):
            shmid = cls.shmget(cls.IPC_PRIVATE, 4 * 1024 * 1024, 0o600)
            if shmid <= 0:
                return False
            cls.shmctl(shmid, cls.IPC_RMID, 0)
            return True

        def __init__(self, size, shmid=None):
            self.size = size
            self.shmid = None
            self.ptr = None
            if shmid is None:
                self.is_creator = True
                self.shmid = self.shmget(self.IPC_PRIVATE, self.size, 0o600)
                self.check_ccall_error(self.shmid > 0, "shmget")
                print("SharedMem[pid %i]: New shmid: %i (size %i)" % (os.getpid(), self.shmid, self.size))
                import atexit

                atexit.register(self.remove)
            else:
                self.is_creator = False
                self.shmid = shmid
                assert self.shmid > 0
            self.ptr = self.shmat(self.shmid, 0, 0)
            self.check_ccall_error(self.ptr != self.ctypes.c_void_p(-1).value, "shmat")
            self.check_ccall_error(self.ptr > 0, "shmat")

        def remove(self):
            if self.ptr:
                self.shmdt(self.ptr)
                self.ptr = None
            if self.shmid and self.shmid > 0:
                if self.is_creator:
                    print("SharedMem[pid %i]: Removing shmid %i (size %i)" % (os.getpid(), self.shmid, self.size))
                    self.shmctl(self.shmid, self.IPC_RMID, 0)
                self.shmid = None

        def __del__(self):
            self.remove()

        def __getstate__(self):
            return {"size": self.size, "shmid": self.shmid}

        def __setstate__(self, state):
            self.__init__(**state)

        def __repr__(self):
            return "<SharedMem shmid=%r size=%r is_creator=%r>" % (self.shmid, self.size, self.is_creator)


def next_power_of_two(n):
    return 2 ** (int(n - 1).bit_length())


class SharedNumpyArray:
    """
    This class provides a way to create Numpy arrays in shared memory.
    It adds some logic to mark whether some shared memory segment can be reused
    - that is when the client marks it as unused.

    Note that there are a few similar Python modules:
      https://pypi.python.org/pypi/SharedArray
      https://parad0x.org/git/python/shared-array/about
      https://bitbucket.org/cleemesser/numpy-sharedmem/src
      https://stackoverflow.com/questions/5033799/how-do-i-pass-large-numpy-arrays
      https://stackoverflow.com/questions/7894791/use-numpy-array-in-shared-memory
    """

    # cls members
    ServerLock = Lock()
    ServerInstances = set()
    ServerArrayId = 0

    class TooMuchInstances(SharedMem.ShmException):
        pass

    ExtraSpaceBytes = 4096
    # local members
    is_server = False
    mem = None
    shape, strides, typestr = None, None, None

    @staticmethod
    def numpy_strides_for_fortran(shape, typestr):
        itemsize = numpy.dtype(typestr).itemsize
        strides = [itemsize]
        for s in shape:
            strides += [strides[-1] * s]
        strides = strides[:-1]
        return tuple(strides)

    @staticmethod
    def numpy_strides_for_c_contiguous(shape, typestr):
        itemsize = numpy.dtype(typestr).itemsize
        strides = [numpy.prod(shape[i + 1 :], dtype="int") * itemsize for i in range(len(shape))]
        return tuple(strides)

    @classmethod
    def needed_mem_size(cls, shape, typestr):
        itemsize = numpy.dtype(typestr).itemsize
        mem_size = cls.ExtraSpaceBytes + itemsize * numpy.prod(shape)
        return mem_size

    @classmethod
    def as_shared(cls, array):
        assert isinstance(array, numpy.ndarray)
        if isinstance(array.base, SharedNumpyArray):
            assert array.base.is_in_use()
            return array.base
        return cls.create_copy(array)

    @classmethod
    def create_copy(cls, array):
        assert isinstance(array, numpy.ndarray)
        array_intf = array.__array_interface__
        shape = array_intf["shape"]
        strides = array_intf["strides"]
        typestr = array_intf["typestr"]
        if array.flags.c_contiguous or array.flags.f_contiguous:
            pass  # ok, we can reuse it like that
        else:
            assert strides
            # Use some similar strides so that the copying might be faster.
            if strides[0] == array.itemsize:
                strides = cls.numpy_strides_for_fortran(shape=shape, typestr=typestr)
            else:
                strides = None  # C-contiguous
        inst = cls.create_new(shape=shape, strides=strides, typestr=typestr)
        inst.create_numpy_array()[...] = array
        assert inst._get_sanity_check_flag_ref().value == 42
        assert inst.is_in_use()
        return inst

    @classmethod
    def create_new(cls, shape, strides, typestr):
        needed_mem_size = cls.needed_mem_size(shape=shape, typestr=typestr)
        with cls.ServerLock:
            for inst in cls.ServerInstances:
                assert isinstance(inst, SharedNumpyArray)
                assert inst._get_sanity_check_flag_ref().value == 42
                if inst.is_in_use():
                    continue
                if inst.mem.size < needed_mem_size:
                    inst._init_mem(shape=shape, typestr=typestr)
                # We can reuse it.
                inst._set_new_array_id()
                inst._set_is_used(1)
                inst._set_numpy_format(shape=shape, strides=strides, typestr=typestr)
                return inst
        return cls(shape=shape, strides=strides, typestr=typestr)

    @classmethod
    def _get_new_array_id(cls):
        array_id = cls.ServerArrayId
        cls.ServerArrayId += 1
        return array_id

    def _set_new_array_id(self):
        assert self.is_server
        self.array_id = self._get_new_array_id()
        self._get_array_id_ref().value = self.array_id

    def __init__(self, shape, strides, typestr, mem=None, array_id=None):
        if not mem:
            assert array_id is None
            if len(self.ServerInstances) >= SharedMemNumpyConfig["max_server_instances"]:
                raise self.TooMuchInstances("too much instances (%i)" % len(self.ServerInstances))
            self.is_server = True
            self._init_mem(shape=shape, typestr=typestr)
            self._set_new_array_id()
            self._set_is_used(1)
        else:
            assert array_id is not None
            self.is_server = False
            self.array_id = array_id
            mem_size = self.needed_mem_size(shape=shape, typestr=typestr)
            assert isinstance(mem, SharedMem)
            assert mem.size >= mem_size
            assert mem.shmid > 0
            assert mem.ptr > 0
            self.mem = mem
            assert self._get_sanity_check_flag_ref().value == 42
            assert self._get_array_id_ref().value == self.array_id
            assert self.is_in_use()
        self._set_numpy_format(shape=shape, strides=strides, typestr=typestr)
        if self.is_server:
            with self.ServerLock:
                self.ServerInstances.add(self)

    def _set_numpy_format(self, shape, strides, typestr):
        itemsize = numpy.dtype(typestr).itemsize
        if strides:
            assert all([st > 0 for st in strides])
            assert sum([st * (sh - 1) for (st, sh) in zip(strides, shape)]) + itemsize == numpy.prod(shape) * itemsize
        self.shape = shape
        self.strides = strides
        self.typestr = typestr

    def _init_mem(self, shape, typestr):
        assert self.is_server
        if self.mem:
            self.mem.remove()
            self.mem = None
        assert numpy.prod(shape) > 0
        mem_size = next_power_of_two(self.needed_mem_size(shape=shape, typestr=typestr))
        mem_size = max(SharedMemNumpyConfig["min_shared_mem_size"], mem_size)
        self.mem = SharedMem(size=mem_size)
        self._get_sanity_check_flag_ref().value = 42

    def get_numpy_array_data_ptr(self):
        assert self.mem.ptr > 0
        return self.mem.ptr + self.ExtraSpaceBytes

    @property
    def __array_interface__(self):
        assert self.shape
        # https://docs.scipy.org/doc/numpy/reference/arrays.interface.html
        return {
            "data": (self.get_numpy_array_data_ptr(), False),
            "shape": self.shape,
            "strides": self.strides,
            "typestr": self.typestr,
            "version": 3,
        }

    def create_numpy_array(self):
        assert self._get_sanity_check_flag_ref().value == 42
        assert self._get_array_id_ref().value == self.array_id
        assert self.is_in_use()
        a = numpy.array(self, copy=False)
        assert a.__array_interface__["data"][0] == self.get_numpy_array_data_ptr()
        assert not a.flags.owndata, "a.__array_interface__ = %r" % a.__array_interface__
        assert a.base is self
        assert a.nbytes + self.ExtraSpaceBytes <= self.mem.size
        assert (
            sum([st * (sh - 1) for (st, sh) in zip(a.strides, a.shape)]) + a.itemsize
            == numpy.prod(a.shape) * a.itemsize
            == a.nbytes
        )
        return a

    def _get_sanity_check_flag_ref(self):
        assert self.mem.ptr > 0
        import ctypes

        return ctypes.cast(ctypes.c_void_p(self.mem.ptr), ctypes.POINTER(ctypes.c_uint64)).contents

    def _get_array_id_ref(self):
        assert self.mem.ptr > 0
        import ctypes

        return ctypes.cast(ctypes.c_void_p(self.mem.ptr + 8), ctypes.POINTER(ctypes.c_uint64)).contents

    def _get_in_use_flag_ref(self):
        assert self.mem.ptr > 0
        import ctypes

        return ctypes.cast(ctypes.c_void_p(self.mem.ptr + 16), ctypes.POINTER(ctypes.c_uint64)).contents

    def _set_is_used(self, n):
        self._get_in_use_flag_ref().value = n

    def is_in_use(self):
        return self._get_in_use_flag_ref().value > 0

    def set_unused(self):
        if self.is_server:
            return
        if self.mem:
            self._set_is_used(0)
            self.mem.remove()
            self.mem = None

    def __getstate__(self):
        return {
            "shape": self.shape,
            "strides": self.strides,
            "typestr": self.typestr,
            "mem": self.mem,
            "array_id": self.array_id,
        }

    def __setstate__(self, state):
        self.__init__(**state)

    def __del__(self):
        # On the server side, we will get deleted at program end
        # because we are referenced in the global SharedNumpyArray.ServerInstances.
        # On the client side, we will get deleted once we are not used anymore.
        # Note that self.array holds a reference to self.
        self.set_unused()

    def __repr__(self):
        return "<%s is_server=%r state=%r>" % (self.__class__.__name__, self.is_server, self.__getstate__())


def attrChain(base, *attribs, **kwargs):
    default = kwargs.get("default", None)
    obj = base
    for attr in attribs:
        if obj is None:
            return default
        obj = getattr(obj, attr, None)
    if obj is None:
        return default
    return obj


# This is needed in some cases to avoid pickling problems with bounded funcs.
def funcCall(attrChainArgs, args=()):
    f = attrChain(*attrChainArgs)
    return f(*args)


Unpickler = pickle.Unpickler
if PY3:

    def get_func_closure(f):
        return f.__closure__

    # (code, globals[, name[, argdefs[, closure]]])
    def get_func_tuple(f):
        return (
            f.__code__,
            f.__globals__,
            f.__name__,
            f.__defaults__,
            f.__closure__,
        )

else:

    def get_func_closure(f):
        return f.func_closure

    def get_func_tuple(f):
        return (
            f.func_code,
            f.func_globals,
            f.func_name,
            f.func_defaults,
            f.func_closure,
        )


_closure = (lambda x: lambda: x)(0)
# noinspection PyUnresolvedReferences
_cell = get_func_closure(_closure)[0]
CellType = type(_cell)
ModuleType = type(sys)
# noinspection PyUnresolvedReferences
DictType = dict if PY3 else types.DictionaryType

if PY3:

    class BufferType:
        "Dummy"

    def make_buffer(*args):
        assert False

else:
    # noinspection PyUnresolvedReferences
    make_buffer = buffer
    # noinspection PyUnresolvedReferences
    BufferType = types.BufferType

    def bytes(x, *args):
        return str(x)


if PY3:
    _old_style_class = None

    class OldStyleClass:
        "Dummy"

    class _new_style_class:
        pass

    NewStyleClass = type
else:

    class _old_style_class:
        pass

    class _new_style_class(object):
        pass

    OldStyleClass = type(_old_style_class)  # == types.ClassType (classobj)
    NewStyleClass = type(_new_style_class)  # (type)


def makeFuncCell(value):
    return get_func_closure((lambda: value))[0]


def getModuleDict(modname, path=None):
    """
    :param str modname: such that "import <modname>" would work
    :param list[str] path: sys.path
    :return: the dict of the mod
    :rtype: dict[str]
    """
    try:
        mod = import_module(modname)
    except ImportError:
        # Try again with extended sys.path.
        assert path
        for p in path:
            if p not in sys.path:
                sys.path.append(p)
        mod = import_module(modname)
    return mod.__dict__


def getModNameForModDict(obj):
    """
    :type obj: dict
    :rtype: str | None
    :returns The module name or None. It will not return '__main__' in any case
    because that likely will not be the same in the unpickling environment.
    Also see: https://stackoverflow.com/questions/56171796/
    """
    if "__name__" not in obj:
        return None  # this is not a module
    mod_name = obj["__name__"]
    if mod_name == "__main__":
        return None
    if mod_name not in sys.modules:
        return None  # does not look like we have it loaded
    mod = sys.modules[mod_name]
    if mod.__dict__ is obj:
        return mod_name
    return None


def getNormalDict(d):
    """
    :type d: dict[str] | dictproxy
    :rtype: dict[str]
    It also removes getset_descriptor. New-style classes have those.
    """
    r = {}
    for k, v in d.items():
        if isinstance(v, types.GetSetDescriptorType):
            continue
        r[k] = v
    return r


def make_numpy_ndarray_fromstring(s, dtype, shape):
    return numpy.fromstring(s, dtype=dtype).reshape(shape)


SharedMemNumpyConfig = {
    "enabled": False,
    "auto_pickling_min_size": 8 * 1024 * 1024,  # 8MB
    "min_shared_mem_size": 32 * 1024 * 1024,  # 32MB
    "max_server_instances": 10,
}


def use_shared_mem_for_numpy_array(obj):
    assert isinstance(obj, numpy.ndarray)
    if obj.shape == ():  # scalar
        return False  # cannot use shared memory because it will always use its own memory
    if isinstance(obj.base, SharedNumpyArray):
        assert obj.base.is_in_use()
        return True
    if not SharedMemNumpyConfig["enabled"]:
        return False
    return obj.nbytes >= SharedMemNumpyConfig["auto_pickling_min_size"]


def numpy_set_unused(v):
    """
    :param numpy.ndarray v: array which will be marked as not-used-anymore
    This will tell mechanisms like SharedNumpyArray that it can reuse the memory.
    On the client side, this will even unmap the memory, so any further access
    to it will cause a SEGFAULT.
    """
    if v is None:
        return
    assert isinstance(v, numpy.ndarray)
    if isinstance(v.base, SharedNumpyArray):
        assert v.base.is_in_use()  # must not be called multiple times
        v.base.set_unused()


def numpy_copy_and_set_unused(v):
    """
    :param dict[str,numpy.ndarray|object] | numpy.ndarray | object v: object to be handled
    If v is a dict, we will return a new copied dict where every value is mapped through numpy_copy_and_set_unused.
    If v is a numpy.ndarray and its base is a SharedNumpyArray, we will copy it and
      call numpy_set_unused on the old value.
    If v is a numpy.ndarray and its base is not a SharedNumpyArray, we will just return it as it is and do nothing.
    In all other cases, we will also just return the object as it is and do nothing.
    """
    if isinstance(v, numpy.ndarray):
        if isinstance(v.base, SharedNumpyArray):
            newv = v.copy(order="A")
            numpy_set_unused(v)
            return newv
        return v
    if isinstance(v, dict):
        return {k: numpy_copy_and_set_unused(vv) for (k, vv) in v.items()}
    return v


def numpy_alloc(shape, dtype, fortran_for_shared=False):
    """
    If EnableAutoNumpySharedMemPickling is True, this will allocate a Numpy array
    in shared memory so we avoid a copy later on when this Numpy array would
    be transferred to another process via pickling.
    """
    if SharedMemNumpyConfig["enabled"]:
        dtype = numpy.dtype(dtype)
        typestr = dtype.str
        strides = None
        if fortran_for_shared:
            strides = SharedNumpyArray.numpy_strides_for_fortran(shape=shape, typestr=typestr)
        try:
            return SharedNumpyArray.create_new(shape=shape, strides=strides, typestr=typestr)
        except SharedMem.ShmException as e:
            print("numpy_alloc: SharedMem exception: %s" % e)
    # Fallback.
    return numpy.ndarray(shape, dtype=dtype)


try:
    _BasePickler = pickle._Pickler  # use the pure Python implementation
except AttributeError:
    _BasePickler = pickle.Pickler


class Pickler(_BasePickler):
    """
    We extend the standard Pickler to be able to pickle some more types,
    such as lambdas and functions, code, func cells, buffer and more.
    """

    def __init__(self, *args, **kwargs):
        if not "protocol" in kwargs:
            kwargs["protocol"] = pickle.HIGHEST_PROTOCOL
        _BasePickler.__init__(self, *args, **kwargs)

    dispatch = _BasePickler.dispatch.copy()

    def save_func(self, obj):
        try:
            self.save_global(obj)
            return
        except pickle.PicklingError:
            pass
        assert type(obj) is types.FunctionType
        self.save(types.FunctionType)
        self.save(get_func_tuple(obj))
        self.write(pickle.REDUCE)
        if id(obj) not in self.memo:  # Could be if we recursively landed here. See also pickle.save_tuple().
            self.memoize(obj)

    dispatch[types.FunctionType] = save_func

    def save_method(self, obj):
        try:
            self.save_global(obj)
            return
        except pickle.PicklingError:
            pass
        assert type(obj) is types.MethodType
        self.save(types.MethodType)
        if PY3:
            self.save((obj.__func__, obj.__self__))
        else:
            self.save((obj.im_func, obj.im_self, obj.im_class))
        self.write(pickle.REDUCE)
        self.memoize(obj)

    dispatch[types.MethodType] = save_method

    def save_code(self, obj):
        assert type(obj) is types.CodeType
        self.save(marshal.loads)
        self.save((marshal.dumps(obj),))
        self.write(pickle.REDUCE)
        self.memoize(obj)

    dispatch[types.CodeType] = save_code

    def save_cell(self, obj):
        assert type(obj) is CellType
        self.save(makeFuncCell)
        self.save((obj.cell_contents,))
        self.write(pickle.REDUCE)
        self.memoize(obj)

    dispatch[CellType] = save_cell

    # We also search for module dicts and reference them.
    # This is for FunctionType.func_globals.
    def intellisave_dict(self, obj):
        modname = getModNameForModDict(obj)
        if modname:
            self.save(getModuleDict)
            self.save((modname, sys.path))
            self.write(pickle.REDUCE)
            self.memoize(obj)
            return
        self.save_dict(obj)

    dispatch[DictType] = intellisave_dict

    def save_module(self, obj):
        modname = getModNameForModDict(obj.__dict__)
        if modname:
            self.save(import_module)
            self.save((modname,))
            self.write(pickle.REDUCE)
            self.memoize(obj)
            return
        # We could maybe construct it manually. For now, just fail.
        raise pickle.PicklingError("cannot pickle module %r" % obj)

    dispatch[ModuleType] = save_module

    def save_buffer(self, obj):
        self.save(buffer)
        self.save((str(obj),))
        self.write(pickle.REDUCE)

    dispatch[BufferType] = save_buffer

    def save_string(self, obj, pack=struct.pack):
        # Difference to base: We just always use BINSTRING (simpler)
        # and use a separate write for the obj itself.
        # For a huge obj, this avoids one unnecessary copy of the data.
        self.write(pickle.BINSTRING + pack("<i", len(obj)))
        self.write(bytes(obj, "utf8"))

    dispatch[str] = save_string

    def save_ndarray(self, obj):
        if use_shared_mem_for_numpy_array(obj):
            try:
                shared = SharedNumpyArray.as_shared(obj)
            except SharedMem.ShmException as e:
                print("SharedNumpyArray exception: %s" % e)
                # fallback to default
            else:
                self.save(shared.create_numpy_array)
                self.save(())
                self.write(pickle.REDUCE)
                return
        # For some reason, Numpy fromstring/tostring is faster than Numpy loads/dumps.
        self.save(make_numpy_ndarray_fromstring)
        self.save((obj.tostring(), str(obj.dtype), obj.shape))
        self.write(pickle.REDUCE)

    dispatch[numpy.ndarray] = save_ndarray

    def save_iobuffer_dummy(self, obj):
        # Not supported but we want to not fail and just store None.
        self.save_none(None)

    dispatch[io.BufferedReader] = save_iobuffer_dummy
    dispatch[io.BufferedWriter] = save_iobuffer_dummy

    # Overwrite to avoid the broken pickle.whichmodule() which might return "__main__".
    def save_global(self, obj, name=None):
        assert obj
        assert id(obj) not in self.memo
        if name is None:
            name = obj.__name__

        module = getattr(obj, "__module__", None)
        if module == "__main__" and globals().get(name, None) is obj:
            # Can happen if this is directly executed.
            module = __name__  # should be correct now
        if module is None or module == "__main__":
            module = pickle.whichmodule(obj, name)
        if module is None or module == "__main__":
            raise pickle.PicklingError("Can't pickle %r: module not found: %s" % (obj, module))

        try:
            __import__(module)
            mod = sys.modules[module]
            klass = getattr(mod, name)
        except (ImportError, KeyError, AttributeError):
            raise pickle.PicklingError("Can't pickle %r: it's not found as %s.%s" % (obj, module, name))
        else:
            if klass is not obj:
                raise pickle.PicklingError("Can't pickle %r: it's not the same object as %s.%s" % (obj, module, name))

        assert "\n" not in module
        assert "\n" not in name
        self.write(pickle.GLOBAL + bytes(module + "\n" + name + "\n", "utf8"))
        self.memoize(obj)

    def save_type(self, obj):
        try:
            self.save_global(obj)
            return
        except pickle.PicklingError:
            pass
        # Some types in the types modules are not correctly referenced,
        # such as types.FunctionType. This is fixed here.
        for modname in ["types"]:
            moddict = sys.modules[modname].__dict__
            for modobjname, modobj in moddict.items():
                if modobj is obj:
                    self.write(pickle.GLOBAL + bytes(modname + "\n" + modobjname + "\n", "utf8"))
                    self.memoize(obj)
                    return
        # Generic serialization of new-style classes.
        self.save(type)
        self.save((obj.__name__, obj.__bases__, getNormalDict(obj.__dict__)))
        self.write(pickle.REDUCE)
        self.memoize(obj)

    dispatch[NewStyleClass] = save_type

    # This is about old-style classes.
    def save_class(self, cls):
        try:
            # First try with a global reference. This works normally. This is the default original pickle behavior.
            self.save_global(cls)
            return
        except pickle.PicklingError:
            pass
        # It didn't worked. But we can still serialize it.
        # Note that this could potentially confuse the code if the class is reference-able in some other way
        # - then we will end up with two versions of the same class.
        self.save(types.ClassType)
        self.save((cls.__name__, cls.__bases__, cls.__dict__))
        self.write(pickle.REDUCE)
        self.memoize(cls)
        return

    dispatch[OldStyleClass] = save_class

    # avoid pickling instances of ourself. this mostly doesn't make sense and leads to trouble.
    # however, also doesn't break. it mostly makes sense to just ignore.
    def __getstate__(self):
        return None

    def __setstate__(self, state):
        pass
