# exchangev2.py - repository exchange for wire protocol version 2

#

# Copyright 2018 Gregory Szorc <gregory.szorc@gmail.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 collections

import weakref

from .i18n import _

from .node import (

    nullid,

    short,

)

from . import (

    bookmarks,

    error,

    mdiff,

    narrowspec,

    phases,

    pycompat,

    setdiscovery,

)

from .interfaces import repository

def pull(pullop):

    """Pull using wire protocol version 2."""

    repo = pullop.repo

    remote = pullop.remote

    usingrawchangelogandmanifest = _checkuserawstorefiledata(pullop)

    # If this is a clone and it was requested to perform a "stream clone",

    # we obtain the raw files data from the remote then fall back to an

    # incremental pull. This is somewhat hacky and is not nearly robust enough

    # for long-term usage.

    if usingrawchangelogandmanifest:

        with repo.transaction(b'clone'):

            _fetchrawstorefiles(repo, remote)

            repo.invalidate(clearfilecache=True)

    tr = pullop.trmanager.transaction()

    # We don't use the repo's narrow matcher here because the patterns passed

    # to exchange.pull() could be different.

    narrowmatcher = narrowspec.match(

        repo.root,

        # Empty maps to nevermatcher. So always

        # set includes if missing.

        pullop.includepats or {b'path:.'},

        pullop.excludepats,

    )

    if pullop.includepats or pullop.excludepats:

        pathfilter = {}

        if pullop.includepats:

            pathfilter[b'include'] = sorted(pullop.includepats)

        if pullop.excludepats:

            pathfilter[b'exclude'] = sorted(pullop.excludepats)

    else:

        pathfilter = None

    # Figure out what needs to be fetched.

    common, fetch, remoteheads = _pullchangesetdiscovery(

        repo, remote, pullop.heads, abortwhenunrelated=pullop.force

    )

    # And fetch the data.

    pullheads = pullop.heads or remoteheads

    csetres = _fetchchangesets(repo, tr, remote, common, fetch, pullheads)

    # New revisions are written to the changelog. But all other updates

    # are deferred. Do those now.

    # Ensure all new changesets are draft by default. If the repo is

    # publishing, the phase will be adjusted by the loop below.

    if csetres[b'added']:

        phases.registernew(repo, tr, phases.draft, csetres[b'added'])

    # And adjust the phase of all changesets accordingly.

    for phasenumber, phase in phases.phasenames.items():

        if phase == b'secret' or not csetres[b'nodesbyphase'][phase]:

            continue

        phases.advanceboundary(

            repo, tr, phasenumber, csetres[b'nodesbyphase'][phase],

        )

    # Write bookmark updates.

    bookmarks.updatefromremote(

        repo.ui,

        repo,

        csetres[b'bookmarks'],

        remote.url(),

        pullop.gettransaction,

        explicit=pullop.explicitbookmarks,

    )

    manres = _fetchmanifests(repo, tr, remote, csetres[b'manifestnodes'])

    # We don't properly support shallow changeset and manifest yet. So we apply

    # depth limiting locally.

    if pullop.depth:

        relevantcsetnodes = set()

        clnode = repo.changelog.node

        for rev in repo.revs(

            b'ancestors(%ln, %s)', pullheads, pullop.depth - 1

        ):

            relevantcsetnodes.add(clnode(rev))

        csetrelevantfilter = lambda n: n in relevantcsetnodes

    else:

        csetrelevantfilter = lambda n: True

    # If obtaining the raw store files, we need to scan the full repo to

    # derive all the changesets, manifests, and linkrevs.

    if usingrawchangelogandmanifest:

        csetsforfiles = []

        mnodesforfiles = []

        manifestlinkrevs = {}

        for rev in repo:

            ctx = repo[rev]

            node = ctx.node()

            if not csetrelevantfilter(node):

                continue

            mnode = ctx.manifestnode()

            csetsforfiles.append(node)

            mnodesforfiles.append(mnode)

            manifestlinkrevs[mnode] = rev

    else:

        csetsforfiles = [n for n in csetres[b'added'] if csetrelevantfilter(n)]

        mnodesforfiles = manres[b'added']

        manifestlinkrevs = manres[b'linkrevs']

    # Find all file nodes referenced by added manifests and fetch those

    # revisions.

    fnodes = _derivefilesfrommanifests(repo, narrowmatcher, mnodesforfiles)

    _fetchfilesfromcsets(

        repo,

        tr,

        remote,

        pathfilter,

        fnodes,

        csetsforfiles,

        manifestlinkrevs,

        shallow=bool(pullop.depth),

    )

def _checkuserawstorefiledata(pullop):

    """Check whether we should use rawstorefiledata command to retrieve data."""

    repo = pullop.repo

    remote = pullop.remote

    # Command to obtain raw store data isn't available.

    if b'rawstorefiledata' not in remote.apidescriptor[b'commands']:

        return False

    # Only honor if user requested stream clone operation.

    if not pullop.streamclonerequested:

        return False

    # Only works on empty repos.

    if len(repo):

        return False

    # TODO This is super hacky. There needs to be a storage API for this. We

    # also need to check for compatibility with the remote.

    if b'revlogv1' not in repo.requirements:

        return False

    return True

def _fetchrawstorefiles(repo, remote):

    with remote.commandexecutor() as e:

        objs = e.callcommand(

            b'rawstorefiledata', {b'files': [b'changelog', b'manifestlog'],}

        ).result()

        # First object is a summary of files data that follows.

        overall = next(objs)

        progress = repo.ui.makeprogress(

            _(b'clone'), total=overall[b'totalsize'], unit=_(b'bytes')

        )

        with progress:

            progress.update(0)

            # Next are pairs of file metadata, data.

            while True:

                try:

                    filemeta = next(objs)

                except StopIteration:

                    break

                for k in (b'location', b'path', b'size'):

                    if k not in filemeta:

                        raise error.Abort(

                            _(b'remote file data missing key: %s') % k

                        )

                if filemeta[b'location'] == b'store':

                    vfs = repo.svfs

                else:

                    raise error.Abort(

                        _(b'invalid location for raw file data: %s')

                        % filemeta[b'location']

                    )

                bytesremaining = filemeta[b'size']

                with vfs.open(filemeta[b'path'], b'wb') as fh:

                    while True:

                        try:

                            chunk = next(objs)

                        except StopIteration:

                            break

                        bytesremaining -= len(chunk)

                        if bytesremaining < 0:

                            raise error.Abort(

                                _(

                                    b'received invalid number of bytes for file '

                                    b'data; expected %d, got extra'

                                )

                                % filemeta[b'size']

                            )

                        progress.increment(step=len(chunk))

                        fh.write(chunk)

                        try:

                            if chunk.islast:

                                break

                        except AttributeError:

                            raise error.Abort(

                                _(

                                    b'did not receive indefinite length bytestring '

                                    b'for file data'

                                )

                            )

                if bytesremaining:

                    raise error.Abort(

                        _(

                            b'received invalid number of bytes for'

                            b'file data; expected %d got %d'

                        )

                        % (

                            filemeta[b'size'],

                            filemeta[b'size'] - bytesremaining,

                        )

                    )

def _pullchangesetdiscovery(repo, remote, heads, abortwhenunrelated=True):

    """Determine which changesets need to be pulled."""

    if heads:

        knownnode = repo.changelog.hasnode

        if all(knownnode(head) for head in heads):

            return heads, False, heads

    # TODO wire protocol version 2 is capable of more efficient discovery

    # than setdiscovery. Consider implementing something better.

    common, fetch, remoteheads = setdiscovery.findcommonheads(

        repo.ui, repo, remote, abortwhenunrelated=abortwhenunrelated

    )

    common = set(common)

    remoteheads = set(remoteheads)

    # If a remote head is filtered locally, put it back in the common set.

    # See the comment in exchange._pulldiscoverychangegroup() for more.

    if fetch and remoteheads:

        has_node = repo.unfiltered().changelog.index.has_node

        common |= {head for head in remoteheads if has_node(head)}

        if set(remoteheads).issubset(common):

            fetch = []

    common.discard(nullid)

    return common, fetch, remoteheads

def _fetchchangesets(repo, tr, remote, common, fetch, remoteheads):

    # TODO consider adding a step here where we obtain the DAG shape first

    # (or ask the server to slice changesets into chunks for us) so that

    # we can perform multiple fetches in batches. This will facilitate

    # resuming interrupted clones, higher server-side cache hit rates due

    # to smaller segments, etc.

    with remote.commandexecutor() as e:

        objs = e.callcommand(

            b'changesetdata',

            {

                b'revisions': [

                    {

                        b'type': b'changesetdagrange',

                        b'roots': sorted(common),

                        b'heads': sorted(remoteheads),

                    }

                ],

                b'fields': {b'bookmarks', b'parents', b'phase', b'revision'},

            },

        ).result()

        # The context manager waits on all response data when exiting. So

        # we need to remain in the context manager in order to stream data.

        return _processchangesetdata(repo, tr, objs)

def _processchangesetdata(repo, tr, objs):

    repo.hook(b'prechangegroup', throw=True, **pycompat.strkwargs(tr.hookargs))

    urepo = repo.unfiltered()

    cl = urepo.changelog

    cl.delayupdate(tr)

    # The first emitted object is a header describing the data that

    # follows.

    meta = next(objs)

    progress = repo.ui.makeprogress(

        _(b'changesets'), unit=_(b'chunks'), total=meta.get(b'totalitems')

    )

    manifestnodes = {}

    def linkrev(node):

        repo.ui.debug(b'add changeset %s\n' % short(node))

        # Linkrev for changelog is always self.

        return len(cl)

    def onchangeset(cl, node):

        progress.increment()

        revision = cl.changelogrevision(node)

        # We need to preserve the mapping of changelog revision to node

        # so we can set the linkrev accordingly when manifests are added.

        manifestnodes[cl.rev(node)] = revision.manifest

    nodesbyphase = {phase: set() for phase in phases.phasenames.values()}

    remotebookmarks = {}

    # addgroup() expects a 7-tuple describing revisions. This normalizes

    # the wire data to that format.

    #

    # This loop also aggregates non-revision metadata, such as phase

    # data.

    def iterrevisions():

        for cset in objs:

            node = cset[b'node']

            if b'phase' in cset:

                nodesbyphase[cset[b'phase']].add(node)

            for mark in cset.get(b'bookmarks', []):

                remotebookmarks[mark] = node

            # TODO add mechanism for extensions to examine records so they

            # can siphon off custom data fields.

            extrafields = {}

            for field, size in cset.get(b'fieldsfollowing', []):

                extrafields[field] = next(objs)

            # Some entries might only be metadata only updates.

            if b'revision' not in extrafields:

                continue

            data = extrafields[b'revision']

            yield (

                node,

                cset[b'parents'][0],