"""
This class is defined to override standard pickle functionality

The goals of it follow:
-Serialize lambdas and nested functions to compiled byte code
-Deal with main module correctly
-Deal with other non-serializable objects

It does not include an unpickler, as standard python unpickling suffices.

This module was extracted from the `cloud` package, developed by `PiCloud, Inc.
<http://www.picloud.com>`_.

Copyright (c) 2012, Regents of the University of California.
Copyright (c) 2009 `PiCloud, Inc. <http://www.picloud.com>`_.
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the University of California, Berkeley nor the
      names of its contributors may be used to endorse or promote
      products derived from this software without specific prior written
      permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""


import operator
import os
import pickle
import struct
import sys
import types
from functools import partial
import itertools
from copy_reg import _extension_registry, _inverted_registry, _extension_cache
import new
import dis
import traceback

#relevant opcodes
STORE_GLOBAL = chr(dis.opname.index('STORE_GLOBAL'))
DELETE_GLOBAL = chr(dis.opname.index('DELETE_GLOBAL'))
LOAD_GLOBAL = chr(dis.opname.index('LOAD_GLOBAL'))
GLOBAL_OPS = [STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL]

HAVE_ARGUMENT = chr(dis.HAVE_ARGUMENT)
EXTENDED_ARG = chr(dis.EXTENDED_ARG)

import logging
cloudLog = logging.getLogger("Cloud.Transport")

try:
    import ctypes
except (MemoryError, ImportError):
    logging.warning('Exception raised on importing ctypes. Likely python bug.. some functionality will be disabled', exc_info = True)
    ctypes = None
    PyObject_HEAD = None
else:

    # for reading internal structures
    PyObject_HEAD = [
        ('ob_refcnt', ctypes.c_size_t),
        ('ob_type', ctypes.c_void_p),
    ]


try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

# These helper functions were copied from PiCloud's util module.
def islambda(func):
    return getattr(func,'func_name') == '<lambda>'

def xrange_params(xrangeobj):
    """Returns a 3 element tuple describing the xrange start, step, and len
    respectively

    Note: Only guarentees that elements of xrange are the same. parameters may
    be different.
    e.g. xrange(1,1) is interpretted as xrange(0,0); both behave the same
    though w/ iteration
    """

    xrange_len = len(xrangeobj)
    if not xrange_len: #empty
        return (0,1,0)
    start = xrangeobj[0]
    if xrange_len == 1: #one element
        return start, 1, 1
    return (start, xrangeobj[1] - xrangeobj[0], xrange_len)

#debug variables intended for developer use:
printSerialization = False
printMemoization = False

useForcedImports = True #Should I use forced imports for tracking?



class CloudPickler(pickle.Pickler):

    dispatch = pickle.Pickler.dispatch.copy()
    savedForceImports = False
    savedDjangoEnv = False #hack tro transport django environment

    def __init__(self, file, protocol=None, min_size_to_save= 0):
        pickle.Pickler.__init__(self,file,protocol)
        self.modules = set() #set of modules needed to depickle
        self.globals_ref = {}  # map ids to dictionary. used to ensure that functions can share global env

    def dump(self, obj):
        # note: not thread safe
        # minimal side-effects, so not fixing
        recurse_limit = 3000
        base_recurse = sys.getrecursionlimit()
        if base_recurse < recurse_limit:
            sys.setrecursionlimit(recurse_limit)
        self.inject_addons()
        try:
            return pickle.Pickler.dump(self, obj)
        except RuntimeError, e:
            if 'recursion' in e.args[0]:
                msg = """Could not pickle object as excessively deep recursion required.
                Try _fast_serialization=2 or contact PiCloud support"""
                raise pickle.PicklingError(msg)
        finally:
            new_recurse = sys.getrecursionlimit()
            if new_recurse == recurse_limit:
                sys.setrecursionlimit(base_recurse)

    def save_buffer(self, obj):
        """Fallback to save_string"""
        pickle.Pickler.save_string(self,str(obj))
    dispatch[buffer] = save_buffer

    #block broken objects
    def save_unsupported(self, obj, pack=None):
        raise pickle.PicklingError("Cannot pickle objects of type %s" % type(obj))
    dispatch[types.GeneratorType] = save_unsupported

    #python2.6+ supports slice pickling. some py2.5 extensions might as well.  We just test it
    try:
        slice(0,1).__reduce__()
    except TypeError: #can't pickle -
        dispatch[slice] = save_unsupported

    #itertools objects do not pickle!
    for v in itertools.__dict__.values():
        if type(v) is type:
            dispatch[v] = save_unsupported


    def save_dict(self, obj):
        """hack fix
        If the dict is a global, deal with it in a special way
        """
        #print 'saving', obj
        if obj is __builtins__:
            self.save_reduce(_get_module_builtins, (), obj=obj)
        else:
            pickle.Pickler.save_dict(self, obj)
    dispatch[pickle.DictionaryType] = save_dict


    def save_module(self, obj, pack=struct.pack):
        """
        Save a module as an import
        """
        #print 'try save import', obj.__name__
        self.modules.add(obj)
        self.save_reduce(subimport,(obj.__name__,), obj=obj)
    dispatch[types.ModuleType] = save_module    #new type

    def save_codeobject(self, obj, pack=struct.pack):
        """
        Save a code object
        """
        #print 'try to save codeobj: ', obj
        args = (
            obj.co_argcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code,
            obj.co_consts, obj.co_names, obj.co_varnames, obj.co_filename, obj.co_name,
            obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, obj.co_cellvars
        )
        self.save_reduce(types.CodeType, args, obj=obj)
    dispatch[types.CodeType] = save_codeobject    #new type

    def save_function(self, obj, name=None, pack=struct.pack):
        """ Registered with the dispatch to handle all function types.

        Determines what kind of function obj is (e.g. lambda, defined at
        interactive prompt, etc) and handles the pickling appropriately.
        """
        write = self.write

        name = obj.__name__
        modname = pickle.whichmodule(obj, name)
        #print 'which gives %s %s %s' % (modname, obj, name)
        try:
            themodule = sys.modules[modname]
        except KeyError: # eval'd items such as namedtuple give invalid items for their function __module__
            modname = '__main__'

        if modname == '__main__':
            themodule = None

        if themodule:
            self.modules.add(themodule)

        if not self.savedDjangoEnv:
            #hack for django - if we detect the settings module, we transport it
            django_settings = os.environ.get('DJANGO_SETTINGS_MODULE', '')
            if django_settings:
                django_mod = sys.modules.get(django_settings)
                if django_mod:
                    cloudLog.debug('Transporting django settings %s during save of %s', django_mod, name)
                    self.savedDjangoEnv = True
                    self.modules.add(django_mod)
                    write(pickle.MARK)
                    self.save_reduce(django_settings_load, (django_mod.__name__,), obj=django_mod)
                    write(pickle.POP_MARK)


        # if func is lambda, def'ed at prompt, is in main, or is nested, then
        # we'll pickle the actual function object rather than simply saving a
        # reference (as is done in default pickler), via save_function_tuple.
        if islambda(obj) or obj.func_code.co_filename == '<stdin>' or themodule is None:
            #Force server to import modules that have been imported in main
            modList = None
            if themodule is None and not self.savedForceImports:
                mainmod = sys.modules['__main__']
                if useForcedImports and hasattr(mainmod,'___pyc_forcedImports__'):
                    modList = list(mainmod.___pyc_forcedImports__)
                self.savedForceImports = True
            self.save_function_tuple(obj, modList)
            return
        else:   # func is nested
            klass = getattr(themodule, name, None)
            if klass is None or klass is not obj:
                self.save_function_tuple(obj, [themodule])
                return

        if obj.__dict__:
            # essentially save_reduce, but workaround needed to avoid recursion
            self.save(_restore_attr)
            write(pickle.MARK + pickle.GLOBAL + modname + '\n' + name + '\n')
            self.memoize(obj)
            self.save(obj.__dict__)
            write(pickle.TUPLE + pickle.REDUCE)
        else:
            write(pickle.GLOBAL + modname + '\n' + name + '\n')
            self.memoize(obj)
    dispatch[types.FunctionType] = save_function

    def save_function_tuple(self, func, forced_imports):
        """  Pickles an actual func object.

        A func comprises: code, globals, defaults, closure, and dict.  We
        extract and save these, injecting reducing functions at certain points
        to recreate the func object.  Keep in mind that some of these pieces
        can contain a ref to the func itself.  Thus, a naive save on these
        pieces could trigger an infinite loop of save's.  To get around that,
        we first create a skeleton func object using just the code (this is
        safe, since this won't contain a ref to the func), and memoize it as
        soon as it's created.  The other stuff can then be filled in later.
        """
        save = self.save
        write = self.write

        # save the modules (if any)
        if forced_imports:
            write(pickle.MARK)
            save(_modules_to_main)
            #print 'forced imports are', forced_imports

            forced_names = map(lambda m: m.__name__, forced_imports)
            save((forced_names,))