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

import os

import tempfile

from .i18n import _

from .node import (

    bin,

    hex,

    nullid,

)

from . import (

    bundle2,

    changegroup as changegroupmod,

    encoding,

    error,

    exchange,

    peer,

    pushkey as pushkeymod,

    pycompat,

    streamclone,

    util,

)

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 abstractserverproto(object):

    """abstract class that summarizes the protocol API

    Used as reference and documentation.

    """

    def getargs(self, args):

        """return the value for arguments in <args>

        returns a list of values (same order as <args>)"""

        raise NotImplementedError()

    def getfile(self, fp):

        """write the whole content of a file into a file like object

        The file is in the form::

            (<chunk-size>\n<chunk>)+0\n

        chunk size is the ascii version of the int.

        """

        raise NotImplementedError()

    def redirect(self):

        """may setup interception for stdout and stderr

        See also the `restore` method."""

        raise NotImplementedError()

    # If the `redirect` function does install interception, the `restore`

    # function MUST be defined. If interception is not used, this function

    # MUST NOT be defined.

    #

    # left commented here on purpose

    #

    #def restore(self):

    #    """reinstall previous stdout and stderr and return intercepted stdout

    #    """

    #    raise NotImplementedError()

class remotebatch(peer.batcher):

    '''batches the queued calls; uses as few roundtrips as possible'''

    def __init__(self, remote):

        '''remote must support _submitbatch(encbatch) and

        _submitone(op, encargs)'''

        peer.batcher.__init__(self)

        self.remote = remote

    def submit(self):

        req, rsp = [], []

        for name, args, opts, resref in self.calls:

            mtd = getattr(self.remote, name)

            batchablefn = getattr(mtd, 'batchable', None)

            if batchablefn is not None:

                batchable = batchablefn(mtd.im_self, *args, **opts)

                encargsorres, encresref = next(batchable)

                if encresref:

                    req.append((name, encargsorres,))

                    rsp.append((batchable, encresref, resref,))

                else:

                    resref.set(encargsorres)

            else:

                if req:

                    self._submitreq(req, rsp)

                    req, rsp = [], []

                resref.set(mtd(*args, **opts))

        if req:

            self._submitreq(req, rsp)

    def _submitreq(self, req, rsp):

        encresults = self.remote._submitbatch(req)

        for encres, r in zip(encresults, rsp):

            batchable, encresref, resref = r

            encresref.set(encres)

            resref.set(next(batchable))

class remoteiterbatcher(peer.iterbatcher):

    def __init__(self, remote):

        super(remoteiterbatcher, self).__init__()

        self._remote = remote

    def __getattr__(self, name):

        if not getattr(self._remote, name, False):

            raise AttributeError(

                'Attempted to iterbatch 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.

        """

        req, rsp = [], []

        for name, args, opts, resref in self.calls:

            mtd = getattr(self._remote, name)

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

            encargsorres, encresref = next(batchable)

            assert encresref

            req.append((name, encargsorres))

            rsp.append((batchable, encresref))

        if req:

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

        self._rsp = rsp

    def results(self):

        for (batchable, encresref), encres in itertools.izip(

                self._rsp, self._resultiter):

            encresref.set(encres)

            yield next(batchable)

# 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 map(bin, 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)

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

             'common': 'nodes',

             'obsmarkers': 'boolean',

             'bundlecaps': 'scsv',

             'listkeys': 'csv',

             'cg': 'boolean',

             'cbattempted': 'boolean'}

# client side

class wirepeer(peer.peerrepository):

    """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.

    """

    def batch(self):

        if self.capable('batch'):

            return remotebatch(self)

        else:

            return peer.localbatch(self)

    def _submitbatch(self, req):

        """run batch request <req> on the server

        Returns an iterator of the raw responses from the server.

        """

        rsp = self._callstream("batch", cmds=encodebatchcmds(req))

        chunk = rsp.read(1024)

        work = [chunk]

        while chunk:

            while ';' not in chunk and chunk:

                chunk = rsp.read(1024)

                work.append(chunk)

            merged = ''.join(work)

            while ';' in merged:

                one, merged = merged.split(';', 1)

                yield unescapearg(one)

            chunk = rsp.read(1024)

            work = [merged, chunk]

        yield unescapearg(''.join(work))

    def _submitone(self, op, args):

        return self._call(op, **args)

    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)

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

    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 = []

        for i in xrange(0, len(pairs), batch):

            n = " ".join([encodelist(p, '-') for p in pairs[i:i + batch]])

            d = self._call("between", pairs=n)

            try:

                r.extend(l and decodelist(l) or [] for l in d.splitlines())

            except ValueError:

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

        return r

    @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

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

    def stream_out(self):

        return self._callstream('stream_out')

    def changegroup(self, nodes, kind):

        n = encodelist(nodes)

        f = self._callcompressable("changegroup", roots=n)

        return changegroupmod.cg1unpacker(f, 'UN')

    def changegroupsubset(self, bases, heads, kind):

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

        bases = encodelist(bases)

        heads = encodelist(heads)

        f = self._callcompressable("changegroupsubset",

                                   bases=bases, heads=heads)