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

    repository,

    setdiscovery,

)

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('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 {'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['added']:

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

    # And adjust the phase of all changesets accordingly.

    for phase in phases.phasenames:

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

            continue

        phases.advanceboundary(repo, tr, phases.phasenames.index(phase),

                               csetres['nodesbyphase'][phase])

    # Write bookmark updates.

    bookmarks.updatefromremote(repo.ui, repo, csetres['bookmarks'],

                               remote.url(), pullop.gettransaction,

                               explicit=pullop.explicitbookmarks)

    manres = _fetchmanifests(repo, tr, remote, csetres['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['added'] if csetrelevantfilter(n)]

        mnodesforfiles = manres['added']

        manifestlinkrevs = manres['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(_('clone'), total=overall[b'totalsize'],

                                        unit=_('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: '

                                        b'%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:

        nodemap = repo.unfiltered().changelog.nodemap

        common |= {head for head in remoteheads if head in nodemap}

        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('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(_('changesets'),

                                    unit=_('chunks'),

                                    total=meta.get(b'totalitems'))

    manifestnodes = {}

    def linkrev(node):

        repo.ui.debug('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}

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

                cset[b'parents'][1],

                # Linknode is always itself for changesets.

                cset[b'node'],

                # We always send full revisions. So delta base is not set.

                nullid,

                mdiff.trivialdiffheader(len(data)) + data,

                # Flags not yet supported.

                0,

            )

    added = cl.addgroup(iterrevisions(), linkrev, weakref.proxy(tr),

                        addrevisioncb=onchangeset)

    progress.complete()

    return {

        'added': added,

        'nodesbyphase': nodesbyphase,

        'bookmarks': remotebookmarks,

        'manifestnodes': manifestnodes,

    }

def _fetchmanifests(repo, tr, remote, manifestnodes):

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

    # Some manifests can be shared between changesets. Filter out revisions

    # we already know about.

    fetchnodes = []

    linkrevs = {}

    seen = set()

    for clrev, node in sorted(manifestnodes.iteritems()):

        if node in seen:

            continue

        try:

            rootmanifest.rev(node)

        except error.LookupError:

            fetchnodes.append(node)

            linkrevs[node] = clrev

        seen.add(node)

    # TODO handle tree manifests

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

    # the wire data to that format.