# changegroup.py - Mercurial changegroup manipulation functions

#

#  Copyright 2006 Olivia Mackall <olivia@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.

import os

import struct

import weakref

from .i18n import _

from .node import (

    hex,

    nullrev,

    short,

)

from .pycompat import open

from . import (

    error,

    match as matchmod,

    mdiff,

    phases,

    pycompat,

    requirements,

    scmutil,

    util,

)

from .interfaces import repository

from .revlogutils import sidedata as sidedatamod

from .revlogutils import constants as revlog_constants

from .utils import storageutil

_CHANGEGROUPV1_DELTA_HEADER = struct.Struct(b"20s20s20s20s")

_CHANGEGROUPV2_DELTA_HEADER = struct.Struct(b"20s20s20s20s20s")

_CHANGEGROUPV3_DELTA_HEADER = struct.Struct(b">20s20s20s20s20sH")

_CHANGEGROUPV4_DELTA_HEADER = struct.Struct(b">B20s20s20s20s20sH")

LFS_REQUIREMENT = b'lfs'

readexactly = util.readexactly

def getchunk(stream):

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

    d = readexactly(stream, 4)

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

    if l <= 4:

        if l:

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

        return b""

    return readexactly(stream, l - 4)

def chunkheader(length):

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

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

def closechunk():

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

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

def _fileheader(path):

    """Obtain a changegroup chunk header for a named path."""

    return chunkheader(len(path)) + path

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, b"wb")

            else:

                # Increase default buffer size because default is usually

                # small (4k is common on Linux).

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

        else:

            fd, filename = pycompat.mkstemp(prefix=b"hg-bundle-", suffix=b".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)

class cg1unpacker:

    """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 = deltaheader.size

    version = b'01'

    _grouplistcount = 1  # One list of files after the manifests

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

        if alg is None:

            alg = b'UN'

        if alg not in util.compengines.supportedbundletypes:

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

        if alg == b'BZ':

            alg = b'_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 != b'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(b">l", d)[0]

        if l <= 4:

            if l:

                raise error.Abort(_(b"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 {b'filename': fname}

    def _deltaheader(self, headertuple, prevnode):

        node, p1, p2, cs = headertuple

        if prevnode is None:

            deltabase = p1

        else:

            deltabase = prevnode

        flags = 0

        protocol_flags = 0

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

    def deltachunk(self, prevnode):

        # Chunkdata: (node, p1, p2, cs, deltabase, delta, flags, sidedata, proto_flags)

        l = self._chunklength()

        if not l:

            return {}

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

        header = self.deltaheader.unpack(headerdata)

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

        header = self._deltaheader(header, prevnode)

        node, p1, p2, deltabase, cs, flags, protocol_flags = header

        return node, p1, p2, cs, deltabase, delta, flags, {}, protocol_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, addrevisioncb=None):

        self.callback = prog.increment

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

        deltas = self.deltaiter()

        storage = repo.manifestlog.getstorage(b'')

        storage.addgroup(deltas, revmap, trp, addrevisioncb=addrevisioncb)

        prog.complete()

        self.callback = None

    def apply(

        self,

        repo,

        tr,

        srctype,

        url,

        targetphase=phases.draft,

        expectedtotal=None,

        sidedata_categories=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

        `sidedata_categories` is an optional set of the remote's sidedata wanted

        categories.

        """

        repo = repo.unfiltered()

        # Only useful if we're adding sidedata categories. If both peers have

        # the same categories, then we simply don't do anything.

        adding_sidedata = (

            (

                requirements.REVLOGV2_REQUIREMENT in repo.requirements

                or requirements.CHANGELOGV2_REQUIREMENT in repo.requirements

            )

            and self.version == b'04'

            and srctype == b'pull'

        )

        if adding_sidedata:

            sidedata_helpers = sidedatamod.get_sidedata_helpers(

                repo,

                sidedata_categories or set(),

                pull=True,

            )

        else:

            sidedata_helpers = None

        def csmap(x):

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

            return len(cl)

        def revmap(x):

            return cl.rev(x)

        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(b'source', srctype)

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

            repo.hook(

                b'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(_(b"adding changesets\n"))

            clstart = len(cl)

            progress = repo.ui.makeprogress(

                _(b'changesets'), unit=_(b'chunks'), total=expectedtotal

            )

            self.callback = progress.increment

            efilesset = set()

            duprevs = []

            def ondupchangelog(cl, rev):

                if rev < clstart:

                    duprevs.append(rev)  # pytype: disable=attribute-error

            def onchangelog(cl, rev):

                ctx = cl.changelogrevision(rev)

                assert efilesset is not None  # help pytype

                efilesset.update(ctx.files)

                repo.register_changeset(rev, ctx)

            self.changelogheader()

            deltas = self.deltaiter()

            if not cl.addgroup(

                deltas,

                csmap,

                trp,

                alwayscache=True,

                addrevisioncb=onchangelog,

                duplicaterevisioncb=ondupchangelog,

            ):

                repo.ui.develwarn(

                    b'applied empty changelog from changegroup',

                    config=b'warn-empty-changegroup',

                )

            efiles = len(efilesset)

            clend = len(cl)

            changesets = clend - clstart

            progress.complete()

            del deltas

            # TODO Python 2.7 removal

            # del efilesset

            efilesset = None

            self.callback = None

            # Keep track of the (non-changelog) revlogs we've updated and their

            # range of new revisions for sidedata rewrite.

            # TODO do something more efficient than keeping the reference to

            # the revlogs, especially memory-wise.

            touched_manifests = {}

            touched_filelogs = {}

            # pull off the manifest group

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