# wireproto.py - generic wire protocol support functions

#

# Copyright 2005-2010 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 hashlib

import os

import tempfile

from .i18n import _

from .node import (

    bin,

    hex,

    nullid,

)

from . import (

    bundle2,

    changegroup as changegroupmod,

    discovery,

    encoding,

    error,

    exchange,

    peer,

    pushkey as pushkeymod,

    pycompat,

    repository,

    streamclone,

    util,

    wireprototypes,

)

from .utils import (

    procutil,

    stringutil,

)

urlerr = util.urlerr

urlreq = util.urlreq

bundle2requiredmain = _('incompatible Mercurial client; bundle2 required')

bundle2requiredhint = _('see https://www.mercurial-scm.org/wiki/'

                        'IncompatibleClient')

bundle2required = '%s\n(%s)\n' % (bundle2requiredmain, bundle2requiredhint)

class remoteiterbatcher(peer.iterbatcher):

    def __init__(self, remote):

        super(remoteiterbatcher, self).__init__()

        self._remote = remote

    def __getattr__(self, name):

        # Validate this method is batchable, since submit() only supports

        # batchable methods.

        fn = getattr(self._remote, name)

        if not getattr(fn, 'batchable', None):

            raise error.ProgrammingError('Attempted to batch a non-batchable '

                                         'call to %r' % name)

        return super(remoteiterbatcher, self).__getattr__(name)

    def submit(self):

        """Break the batch request into many patch calls and pipeline them.

        This is mostly valuable over http where request sizes can be

        limited, but can be used in other places as well.

        """

        # 2-tuple of (command, arguments) that represents what will be

        # sent over the wire.

        requests = []

        # 4-tuple of (command, final future, @batchable generator, remote

        # future).

        results = []

        for command, args, opts, finalfuture in self.calls:

            mtd = getattr(self._remote, command)

            batchable = mtd.batchable(mtd.__self__, *args, **opts)

            commandargs, fremote = next(batchable)

            assert fremote

            requests.append((command, commandargs))

            results.append((command, finalfuture, batchable, fremote))

        if requests:

            self._resultiter = self._remote._submitbatch(requests)

        self._results = results

    def results(self):

        for command, finalfuture, batchable, remotefuture in self._results:

            # Get the raw result, set it in the remote future, feed it

            # back into the @batchable generator so it can be decoded, and

            # set the result on the final future to this value.

            remoteresult = next(self._resultiter)

            remotefuture.set(remoteresult)

            finalfuture.set(next(batchable))

            # Verify our @batchable generators only emit 2 values.

            try:

                next(batchable)

            except StopIteration:

                pass

            else:

                raise error.ProgrammingError('%s @batchable generator emitted '

                                             'unexpected value count' % command)

            yield finalfuture.value

# Forward a couple of names from peer to make wireproto interactions

# slightly more sensible.

batchable = peer.batchable

future = peer.future

# list of nodes encoding / decoding

def decodelist(l, sep=' '):

    if l:

        return [bin(v) for v in  l.split(sep)]

    return []

def encodelist(l, sep=' '):

    try:

        return sep.join(map(hex, l))

    except TypeError:

        raise

# batched call argument encoding

def escapearg(plain):

    return (plain

            .replace(':', ':c')

            .replace(',', ':o')

            .replace(';', ':s')

            .replace('=', ':e'))

def unescapearg(escaped):

    return (escaped

            .replace(':e', '=')

            .replace(':s', ';')

            .replace(':o', ',')

            .replace(':c', ':'))

def encodebatchcmds(req):

    """Return a ``cmds`` argument value for the ``batch`` command."""

    cmds = []

    for op, argsdict in req:

        # Old servers didn't properly unescape argument names. So prevent

        # the sending of argument names that may not be decoded properly by

        # servers.

        assert all(escapearg(k) == k for k in argsdict)

        args = ','.join('%s=%s' % (escapearg(k), escapearg(v))

                        for k, v in argsdict.iteritems())

        cmds.append('%s %s' % (op, args))

    return ';'.join(cmds)

def clientcompressionsupport(proto):

    """Returns a list of compression methods supported by the client.

    Returns a list of the compression methods supported by the client

    according to the protocol capabilities. If no such capability has

    been announced, fallback to the default of zlib and uncompressed.

    """

    for cap in proto.getprotocaps():

        if cap.startswith('comp='):

            return cap[5:].split(',')

    return ['zlib', 'none']

# mapping of options accepted by getbundle and their types

#

# Meant to be extended by extensions. It is extensions responsibility to ensure

# such options are properly processed in exchange.getbundle.

#

# supported types are:

#

# :nodes: list of binary nodes

# :csv:   list of comma-separated values

# :scsv:  list of comma-separated values return as set

# :plain: string with no transformation needed.

gboptsmap = {'heads':  'nodes',

             'bookmarks': 'boolean',

             'common': 'nodes',

             'obsmarkers': 'boolean',

             'phases': 'boolean',

             'bundlecaps': 'scsv',

             'listkeys': 'csv',

             'cg': 'boolean',

             'cbattempted': 'boolean',

             'stream': 'boolean',

}

# client side

class wirepeer(repository.legacypeer):

    """Client-side interface for communicating with a peer repository.

    Methods commonly call wire protocol commands of the same name.

    See also httppeer.py and sshpeer.py for protocol-specific

    implementations of this interface.

    """

    # Begin of ipeercommands interface.

    def iterbatch(self):

        return remoteiterbatcher(self)

    @batchable

    def lookup(self, key):

        self.requirecap('lookup', _('look up remote revision'))

        f = future()

        yield {'key': encoding.fromlocal(key)}, f

        d = f.value

        success, data = d[:-1].split(" ", 1)

        if int(success):

            yield bin(data)

        else:

            self._abort(error.RepoError(data))

    @batchable

    def heads(self):

        f = future()

        yield {}, f

        d = f.value

        try:

            yield decodelist(d[:-1])

        except ValueError:

            self._abort(error.ResponseError(_("unexpected response:"), d))

    @batchable

    def known(self, nodes):

        f = future()

        yield {'nodes': encodelist(nodes)}, f

        d = f.value

        try:

            yield [bool(int(b)) for b in d]

        except ValueError:

            self._abort(error.ResponseError(_("unexpected response:"), d))

    @batchable

    def branchmap(self):

        f = future()

        yield {}, f

        d = f.value

        try:

            branchmap = {}

            for branchpart in d.splitlines():

                branchname, branchheads = branchpart.split(' ', 1)

                branchname = encoding.tolocal(urlreq.unquote(branchname))

                branchheads = decodelist(branchheads)

                branchmap[branchname] = branchheads

            yield branchmap

        except TypeError:

            self._abort(error.ResponseError(_("unexpected response:"), d))

    @batchable

    def listkeys(self, namespace):

        if not self.capable('pushkey'):

            yield {}, None

        f = future()

        self.ui.debug('preparing listkeys for "%s"\n' % namespace)

        yield {'namespace': encoding.fromlocal(namespace)}, f

        d = f.value

        self.ui.debug('received listkey for "%s": %i bytes\n'

                      % (namespace, len(d)))

        yield pushkeymod.decodekeys(d)

    @batchable

    def pushkey(self, namespace, key, old, new):

        if not self.capable('pushkey'):

            yield False, None

        f = future()

        self.ui.debug('preparing pushkey for "%s:%s"\n' % (namespace, key))

        yield {'namespace': encoding.fromlocal(namespace),

               'key': encoding.fromlocal(key),

               'old': encoding.fromlocal(old),

               'new': encoding.fromlocal(new)}, f

        d = f.value

        d, output = d.split('\n', 1)

        try:

            d = bool(int(d))

        except ValueError:

            raise error.ResponseError(

                _('push failed (unexpected response):'), d)

        for l in output.splitlines(True):

            self.ui.status(_('remote: '), l)

        yield d

    def stream_out(self):

        return self._callstream('stream_out')

    def getbundle(self, source, **kwargs):

        kwargs = pycompat.byteskwargs(kwargs)

        self.requirecap('getbundle', _('look up remote changes'))

        opts = {}

        bundlecaps = kwargs.get('bundlecaps') or set()

        for key, value in kwargs.iteritems():

            if value is None:

                continue

            keytype = gboptsmap.get(key)

            if keytype is None:

                raise error.ProgrammingError(

                    'Unexpectedly None keytype for key %s' % key)

            elif keytype == 'nodes':

                value = encodelist(value)

            elif keytype == 'csv':

                value = ','.join(value)

            elif keytype == 'scsv':

                value = ','.join(sorted(value))

            elif keytype == 'boolean':

                value = '%i' % bool(value)

            elif keytype != 'plain':

                raise KeyError('unknown getbundle option type %s'

                               % keytype)

            opts[key] = value

        f = self._callcompressable("getbundle", **pycompat.strkwargs(opts))

        if any((cap.startswith('HG2') for cap in bundlecaps)):

            return bundle2.getunbundler(self.ui, f)

        else:

            return changegroupmod.cg1unpacker(f, 'UN')

    def unbundle(self, cg, heads, url):

        '''Send cg (a readable file-like object representing the

        changegroup to push, typically a chunkbuffer object) to the

        remote server as a bundle.

        When pushing a bundle10 stream, return an integer indicating the

        result of the push (see changegroup.apply()).

        When pushing a bundle20 stream, return a bundle20 stream.

        `url` is the url the client thinks it's pushing to, which is

        visible to hooks.

        '''

        if heads != ['force'] and self.capable('unbundlehash'):

            heads = encodelist(['hashed',

                                hashlib.sha1(''.join(sorted(heads))).digest()])

        else:

            heads = encodelist(heads)

        if util.safehasattr(cg, 'deltaheader'):

            # this a bundle10, do the old style call sequence

            ret, output = self._callpush("unbundle", cg, heads=heads)

            if ret == "":

                raise error.ResponseError(

                    _('push failed:'), output)

            try:

                ret = int(ret)

            except ValueError:

                raise error.ResponseError(

                    _('push failed (unexpected response):'), ret)

            for l in output.splitlines(True):

                self.ui.status(_('remote: '), l)

        else:

            # bundle2 push. Send a stream, fetch a stream.

            stream = self._calltwowaystream('unbundle', cg, heads=heads)

            ret = bundle2.getunbundler(self.ui, stream)

        return ret

    # End of ipeercommands interface.

    # Begin of ipeerlegacycommands interface.

    def branches(self, nodes):

        n = encodelist(nodes)

        d = self._call("branches", nodes=n)

        try:

            br = [tuple(decodelist(b)) for b in d.splitlines()]

            return br

        except ValueError:

            self._abort(error.ResponseError(_("unexpected response:"), d))

    def between(self, pairs):

        batch = 8 # avoid giant requests

        r = []