# changegroup.py - Mercurial changegroup manipulation functions

#

#  Copyright 2006 Matt Mackall <mpm@selenic.com>

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

from __future__ import absolute_import

import os

import struct

import tempfile

import weakref

from .i18n import _

from .node import (

    hex,

    nullrev,

    short,

)

from . import (

    dagutil,

    discovery,

    error,

    mdiff,

    phases,

    pycompat,

    util,

)

_CHANGEGROUPV1_DELTA_HEADER = "20s20s20s20s"

_CHANGEGROUPV2_DELTA_HEADER = "20s20s20s20s20s"

_CHANGEGROUPV3_DELTA_HEADER = ">20s20s20s20s20sH"

def readexactly(stream, n):

    '''read n bytes from stream.read and abort if less was available'''

    s = stream.read(n)

    if len(s) < n:

        raise error.Abort(_("stream ended unexpectedly"

                           " (got %d bytes, expected %d)")

                          % (len(s), n))

    return s

def getchunk(stream):

    """return the next chunk from stream as a string"""

    d = readexactly(stream, 4)

    l = struct.unpack(">l", d)[0]

    if l <= 4:

        if l:

            raise error.Abort(_("invalid chunk length %d") % l)

        return ""

    return readexactly(stream, l - 4)

def chunkheader(length):

    """return a changegroup chunk header (string)"""

    return struct.pack(">l", length + 4)

def closechunk():

    """return a changegroup chunk header (string) for a zero-length chunk"""

    return struct.pack(">l", 0)

def writechunks(ui, chunks, filename, vfs=None):

    """Write chunks to a file and return its filename.

    The stream is assumed to be a bundle file.

    Existing files will not be overwritten.

    If no filename is specified, a temporary file is created.

    """

    fh = None

    cleanup = None

    try:

        if filename:

            if vfs:

                fh = vfs.open(filename, "wb")

            else:

                # Increase default buffer size because default is usually

                # small (4k is common on Linux).

                fh = open(filename, "wb", 131072)

        else:

            fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg")

            fh = os.fdopen(fd, pycompat.sysstr("wb"))

        cleanup = filename

        for c in chunks:

            fh.write(c)

        cleanup = None

        return filename

    finally:

        if fh is not None:

            fh.close()

        if cleanup is not None:

            if filename and vfs:

                vfs.unlink(cleanup)

            else:

                os.unlink(cleanup)

class cg1unpacker(object):

    """Unpacker for cg1 changegroup streams.

    A changegroup unpacker handles the framing of the revision data in

    the wire format. Most consumers will want to use the apply()

    method to add the changes from the changegroup to a repository.

    If you're forwarding a changegroup unmodified to another consumer,

    use getchunks(), which returns an iterator of changegroup

    chunks. This is mostly useful for cases where you need to know the

    data stream has ended by observing the end of the changegroup.

    deltachunk() is useful only if you're applying delta data. Most

    consumers should prefer apply() instead.

    A few other public methods exist. Those are used only for

    bundlerepo and some debug commands - their use is discouraged.

    """

    deltaheader = _CHANGEGROUPV1_DELTA_HEADER

    deltaheadersize = struct.calcsize(deltaheader)

    version = '01'

    _grouplistcount = 1 # One list of files after the manifests

    def __init__(self, fh, alg, extras=None):

        if alg is None:

            alg = 'UN'

        if alg not in util.compengines.supportedbundletypes:

            raise error.Abort(_('unknown stream compression type: %s')

                             % alg)

        if alg == 'BZ':

            alg = '_truncatedBZ'

        compengine = util.compengines.forbundletype(alg)

        self._stream = compengine.decompressorreader(fh)

        self._type = alg

        self.extras = extras or {}

        self.callback = None

    # These methods (compressed, read, seek, tell) all appear to only

    # be used by bundlerepo, but it's a little hard to tell.

    def compressed(self):

        return self._type is not None and self._type != 'UN'

    def read(self, l):

        return self._stream.read(l)

    def seek(self, pos):

        return self._stream.seek(pos)

    def tell(self):

        return self._stream.tell()

    def close(self):

        return self._stream.close()

    def _chunklength(self):

        d = readexactly(self._stream, 4)

        l = struct.unpack(">l", d)[0]

        if l <= 4:

            if l:

                raise error.Abort(_("invalid chunk length %d") % l)

            return 0

        if self.callback:

            self.callback()

        return l - 4

    def changelogheader(self):

        """v10 does not have a changelog header chunk"""

        return {}

    def manifestheader(self):

        """v10 does not have a manifest header chunk"""

        return {}

    def filelogheader(self):

        """return the header of the filelogs chunk, v10 only has the filename"""

        l = self._chunklength()

        if not l:

            return {}

        fname = readexactly(self._stream, l)

        return {'filename': fname}

    def _deltaheader(self, headertuple, prevnode):

        node, p1, p2, cs = headertuple

        if prevnode is None:

            deltabase = p1

        else:

            deltabase = prevnode

        flags = 0

        return node, p1, p2, deltabase, cs, flags

    def deltachunk(self, prevnode):

        l = self._chunklength()

        if not l:

            return {}

        headerdata = readexactly(self._stream, self.deltaheadersize)

        header = struct.unpack(self.deltaheader, headerdata)

        delta = readexactly(self._stream, l - self.deltaheadersize)

        node, p1, p2, deltabase, cs, flags = self._deltaheader(header, prevnode)

        return {'node': node, 'p1': p1, 'p2': p2, 'cs': cs,

                'deltabase': deltabase, 'delta': delta, 'flags': flags}

    def getchunks(self):

        """returns all the chunks contains in the bundle

        Used when you need to forward the binary stream to a file or another

        network API. To do so, it parse the changegroup data, otherwise it will

        block in case of sshrepo because it don't know the end of the stream.

        """

        # For changegroup 1 and 2, we expect 3 parts: changelog, manifestlog,

        # and a list of filelogs. For changegroup 3, we expect 4 parts:

        # changelog, manifestlog, a list of tree manifestlogs, and a list of

        # filelogs.

        #

        # Changelog and manifestlog parts are terminated with empty chunks. The

        # tree and file parts are a list of entry sections. Each entry section

        # is a series of chunks terminating in an empty chunk. The list of these

        # entry sections is terminated in yet another empty chunk, so we know

        # we've reached the end of the tree/file list when we reach an empty

        # chunk that was proceeded by no non-empty chunks.

        parts = 0

        while parts < 2 + self._grouplistcount:

            noentries = True

            while True:

                chunk = getchunk(self)

                if not chunk:

                    # The first two empty chunks represent the end of the

                    # changelog and the manifestlog portions. The remaining

                    # empty chunks represent either A) the end of individual

                    # tree or file entries in the file list, or B) the end of

                    # the entire list. It's the end of the entire list if there

                    # were no entries (i.e. noentries is True).

                    if parts < 2:

                        parts += 1

                    elif noentries:

                        parts += 1

                    break

                noentries = False

                yield chunkheader(len(chunk))

                pos = 0

                while pos < len(chunk):

                    next = pos + 2**20

                    yield chunk[pos:next]

                    pos = next

            yield closechunk()

    def _unpackmanifests(self, repo, revmap, trp, prog, numchanges):

        # We know that we'll never have more manifests than we had

        # changesets.

        self.callback = prog(_('manifests'), numchanges)

        # no need to check for empty manifest group here:

        # if the result of the merge of 1 and 2 is the same in 3 and 4,

        # no new manifest will be created and the manifest group will

        # be empty during the pull

        self.manifestheader()

        repo.manifestlog._revlog.addgroup(self, revmap, trp)

        repo.ui.progress(_('manifests'), None)

        self.callback = None

    def apply(self, repo, tr, srctype, url, targetphase=phases.draft,

              expectedtotal=None):

        """Add the changegroup returned by source.read() to this repo.

        srctype is a string like 'push', 'pull', or 'unbundle'.  url is

        the URL of the repo where this changegroup is coming from.

        Return an integer summarizing the change to this repo:

        - nothing changed or no source: 0

        - more heads than before: 1+added heads (2..n)

        - fewer heads than before: -1-removed heads (-2..-n)

        - number of heads stays the same: 1

        """

        repo = repo.unfiltered()

        def csmap(x):

            repo.ui.debug("add changeset %s\n" % short(x))

            return len(cl)

        def revmap(x):

            return cl.rev(x)

        changesets = files = revisions = 0

        try:

            # The transaction may already carry source information. In this

            # case we use the top level data. We overwrite the argument

            # because we need to use the top level value (if they exist)

            # in this function.

            srctype = tr.hookargs.setdefault('source', srctype)

            url = tr.hookargs.setdefault('url', url)

            repo.hook('prechangegroup',

                      throw=True, **pycompat.strkwargs(tr.hookargs))

            # write changelog data to temp files so concurrent readers

            # will not see an inconsistent view

            cl = repo.changelog

            cl.delayupdate(tr)

            oldheads = set(cl.heads())

            trp = weakref.proxy(tr)

            # pull off the changeset group

            repo.ui.status(_("adding changesets\n"))

            clstart = len(cl)

            class prog(object):

                def __init__(self, step, total):

                    self._step = step

                    self._total = total

                    self._count = 1

                def __call__(self):

                    repo.ui.progress(self._step, self._count, unit=_('chunks'),

                                     total=self._total)

                    self._count += 1

            self.callback = prog(_('changesets'), expectedtotal)

            efiles = set()

            def onchangelog(cl, node):

                efiles.update(cl.readfiles(node))

            self.changelogheader()

            cgnodes = cl.addgroup(self, csmap, trp, addrevisioncb=onchangelog)

            efiles = len(efiles)

            if not cgnodes:

                repo.ui.develwarn('applied empty changegroup',

                                  config='empty-changegroup')

            clend = len(cl)

            changesets = clend - clstart

            repo.ui.progress(_('changesets'), None)

            self.callback = None

            # pull off the manifest group

            repo.ui.status(_("adding manifests\n"))

            self._unpackmanifests(repo, revmap, trp, prog, changesets)

            needfiles = {}

            if repo.ui.configbool('server', 'validate'):

                cl = repo.changelog

                ml = repo.manifestlog

                # validate incoming csets have their manifests

                for cset in xrange(clstart, clend):

                    mfnode = cl.changelogrevision(cset).manifest

                    mfest = ml[mfnode].readdelta()

                    # store file cgnodes we must see

                    for f, n in mfest.iteritems():

                        needfiles.setdefault(f, set()).add(n)

            # process the files

            repo.ui.status(_("adding file changes\n"))

            newrevs, newfiles = _addchangegroupfiles(

                repo, self, revmap, trp, efiles, needfiles)

            revisions += newrevs

            files += newfiles

            deltaheads = 0

            if oldheads:

                heads = cl.heads()

                deltaheads = len(heads) - len(oldheads)

                for h in heads:

                    if h not in oldheads and repo[h].closesbranch():

                        deltaheads -= 1

            htext = ""

            if deltaheads:

                htext = _(" (%+d heads)") % deltaheads

            repo.ui.status(_("added %d changesets"

                             " with %d changes to %d files%s\n")

                             % (changesets, revisions, files, htext))

            repo.invalidatevolatilesets()

            if changesets > 0:

                if 'node' not in tr.hookargs:

                    tr.hookargs['node'] = hex(cl.node(clstart))

                    tr.hookargs['node_last'] = hex(cl.node(clend - 1))

                    hookargs = dict(tr.hookargs)

                else:

                    hookargs = dict(tr.hookargs)

                    hookargs['node'] = hex(cl.node(clstart))

                    hookargs['node_last'] = hex(cl.node(clend - 1))

                repo.hook('pretxnchangegroup',

                          throw=True, **pycompat.strkwargs(hookargs))

            added = [cl.node(r) for r in xrange(clstart, clend)]

            phaseall = None

            if srctype in ('push', 'serve'):

                # Old servers can not push the boundary themselves.

                # New servers won't push the boundary if changeset already

                # exists locally as secret

                #

                # We should not use added here but the list of all change in

                # the bundle