# 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,

    nullid,

    nullrev,

    short,

)

from . import (

    branchmap,

    dagutil,

    discovery,

    error,

    mdiff,

    phases,

    util,

)

_CHANGEGROUPV1_DELTA_HEADER = "20s20s20s20s"

_CHANGEGROUPV2_DELTA_HEADER = "20s20s20s20s20s"

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 combineresults(results):

    """logic to combine 0 or more addchangegroup results into one"""

    changedheads = 0

    result = 1

    for ret in results:

        # If any changegroup result is 0, return 0

        if ret == 0:

            result = 0

            break

        if ret < -1:

            changedheads += ret + 1

        elif ret > 1:

            changedheads += ret - 1

    if changedheads > 0:

        result = 1 + changedheads

    elif changedheads < 0:

        result = -1 + changedheads

    return result

bundletypes = {

    "": ("", None),       # only when using unbundle on ssh and old http servers

                          # since the unification ssh accepts a header but there

                          # is no capability signaling it.

    "HG20": (), # special-cased below

    "HG10UN": ("HG10UN", None),

    "HG10BZ": ("HG10", 'BZ'),

    "HG10GZ": ("HG10GZ", 'GZ'),

}

# hgweb uses this list to communicate its preferred type

bundlepriority = ['HG10GZ', 'HG10BZ', 'HG10UN']

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:

                fh = open(filename, "wb")

        else:

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

            fh = os.fdopen(fd, "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)

def writebundle(ui, cg, filename, bundletype, vfs=None, compression=None):

    """Write a bundle file and return its filename.

    Existing files will not be overwritten.

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

    bz2 compression can be turned off.

    The bundle file will be deleted in case of errors.

    """

    if bundletype == "HG20":

        from . import bundle2

        bundle = bundle2.bundle20(ui)

        bundle.setcompression(compression)

        part = bundle.newpart('changegroup', data=cg.getchunks())

        part.addparam('version', cg.version)

        chunkiter = bundle.getchunks()

    else:

        # compression argument is only for the bundle2 case

        assert compression is None

        if cg.version != '01':

            raise error.Abort(_('old bundle types only supports v1 '

                                'changegroups'))

        header, comp = bundletypes[bundletype]

        if comp not in util.compressors:

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

                              % comp)

        z = util.compressors[comp]()

        subchunkiter = cg.getchunks()

        def chunkiter():

            yield header

            for chunk in subchunkiter:

                yield z.compress(chunk)

            yield z.flush()

        chunkiter = chunkiter()

    # parse the changegroup data, otherwise we will block

    # in case of sshrepo because we don't know the end of the stream

    # an empty chunkgroup is the end of the changegroup

    # a changegroup has at least 2 chunkgroups (changelog and manifest).

    # after that, an empty chunkgroup is the end of the changegroup

    return writechunks(ui, chunkiter, filename, vfs=vfs)

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'

    def __init__(self, fh, alg):

        if alg == 'UN':

            alg = None # get more modern without breaking too much

        if not alg in util.decompressors:

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

                             % alg)

        if alg == 'BZ':

            alg = '_truncatedBZ'

        self._stream = util.decompressors[alg](fh)

        self._type = alg

        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

    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

        return node, p1, p2, deltabase, cs

    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 = self._deltaheader(header, prevnode)

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

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

    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.

        """

        # an empty chunkgroup is the end of the changegroup

        # a changegroup has at least 2 chunkgroups (changelog and manifest).

        # after that, an empty chunkgroup is the end of the changegroup

        empty = False

        count = 0

        while not empty or count <= 2:

            empty = True

            count += 1

            while True:

                chunk = getchunk(self)

                if not chunk:

                    break

                empty = 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.manifest.addgroup(self, revmap, trp)

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

    def apply(self, repo, srctype, url, emptyok=False,

              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()

        wasempty = (len(repo.changelog) == 0)

        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

        tr = repo.transaction("\n".join([srctype, util.hidepassword(url)]))

        try:

            # The transaction could have been created before and already

            # carries 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, **tr.hookargs)

            # write changelog data to temp files so concurrent readers

            # will not see an inconsistent view

            cl = repo.changelog

            cl.delayupdate(tr)

            oldheads = 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.read(node)[3])

            self.changelogheader()

            srccontent = cl.addgroup(self, csmap, trp,

                                     addrevisioncb=onchangelog)

            efiles = len(efiles)

            if not (srccontent or emptyok):

                raise error.Abort(_("received changelog group is empty"))

            clend = len(cl)

            changesets = clend - clstart

            repo.ui.progress(_('changesets'), 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', default=False):

                # validate incoming csets have their manifests

                for cset in xrange(clstart, clend):

                    mfnode = repo.changelog.read(repo.changelog.node(cset))[0]

                    mfest = repo.manifest.readdelta(mfnode)

                    # store file nodes we must see

                    for f, n in mfest.iteritems():

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