"""Support for export of multidatabases."""

##############################################################################
#
# Based on the ZODB import/export code.
# Copyright (c) 2009 David Glick.
# All Rights Reserved.
#
# This software is subject to the provisions of the Zope Public License,
# Version 2.1 (ZPL).  A copy of the ZPL should accompany this distribution.
# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
# FOR A PARTICULAR PURPOSE
#
##############################################################################

import logging
import cPickle, cStringIO
from ZODB.utils import p64, u64
from ZODB.ExportImport import export_end_marker
from ZODB.DemoStorage import DemoStorage

logger = logging.getLogger('multiexport')

def export_zexp(self, fname):
    context = self
    f = open(fname, 'wb')
    f.write('ZEXP')
    for oid, p in flatten_multidatabase(context):
        f.writelines((oid, p64(len(p)), p))
    f.write(export_end_marker)
    f.close()

def flatten_multidatabase(context):
    """Walk a multidatabase and yield rewritten pickles with oids for a single database"""
    base_oid = context._p_oid
    base_conn = context._p_jar
    dbs = base_conn.connections
    
    dummy_storage = DemoStorage()

    oids = [(base_conn._db.database_name, base_oid)]
    done_oids = {}
    # table to keep track of mapping old oids to new oids
    ooid_to_oid = {oids[0]: dummy_storage.new_oid()}
    while oids:
        # loop while references remain to objects we haven't exported yet
        (dbname, ooid) = oids.pop(0)
        if (dbname, ooid) in done_oids:
            continue
        done_oids[(dbname, ooid)] = True

        db = dbs[dbname]
        try:
            # get pickle
            p, serial = db._storage.load(ooid, db._version)
        except:
            logger.debug("broken reference for db %s, oid %s", (dbname, repr(ooid)),
                         exc_info=True)
        else:
            def persistent_load(ref):
                """ Remap a persistent id to a new ID and create a ghost for it.
                
                This is called by the unpickler for each reference found.
                """

                # resolve the reference to a database name and oid
                if isinstance(ref, tuple):
                    rdbname, roid = (dbname, ref[0])
                elif isinstance(ref, str):
                    rdbname, roid = (dbname, ref)
                else:
                    try:
                        ref_type, args = ref
                    except ValueError:
                        # weakref
                        return
                    else:
                        if ref_type in ('m', 'n'):
                            rdbname, roid = (args[0], args[1])
                        else:
                            return

                # traverse Products.ZODBMountpoint mountpoints to the mounted location
                rdb = dbs[rdbname]
                p, serial = rdb._storage.load(roid, rdb._version)
                klass = p.split()[0]
                if 'MountedObject' in klass:
                    mountpoint = rdb.get(roid)
                    # get the object with the root as a parent, then unwrap,
                    # since there's no API to get the unwrapped object
                    mounted = mountpoint._getOrOpenObject(app).aq_base
                    rdbname = mounted._p_jar._db.database_name
                    roid = mounted._p_oid

                if roid:
                    print '%s:%s -> %s:%s' % (dbname, u64(ooid), rdbname, u64(roid))
                    oids.append((rdbname, roid))

                try:
                    oid = ooid_to_oid[(rdbname, roid)]
                except KeyError:
                    # generate a new oid and associate it with this old db/oid
                    ooid_to_oid[(rdbname, roid)] = oid = dummy_storage.new_oid()
                return Ghost(oid)

            # do the repickling dance to rewrite references
            
            pfile = cStringIO.StringIO(p)
            unpickler = cPickle.Unpickler(pfile)
            unpickler.persistent_load = persistent_load

            newp = cStringIO.StringIO()
            pickler = cPickle.Pickler(newp, 1)
            pickler.persistent_id = persistent_id

            pickler.dump(unpickler.load())
            pickler.dump(unpickler.load())
            p = newp.getvalue()

            yield ooid_to_oid[(dbname, ooid)], p

class Ghost(object):
    __slots__ = ("oid",)
    def __init__(self, oid):
        self.oid = oid

def persistent_id(obj):
    if isinstance(obj, Ghost):
        return obj.oid

export_zexp(app.mysite, '/tmp/mysite.zexp')