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

import sys

import tempfile

import urllib

from .i18n import _

from .node import (

    bin,

    hex,

)

from . import (

    bundle2,

    changegroup as changegroupmod,

    encoding,

    error,

    exchange,

    peer,

    pushkey as pushkeymod,

    streamclone,

    util,

)

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

    def groupchunks(self, cg):

        """return 4096 chunks from a changegroup object

        Some protocols may have compressed the contents."""

        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 = batchable.next()

                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(batchable.next())

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

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

        cmds = []

        for op, argsdict in req:

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

                            for k, v in argsdict.iteritems())

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

        rsp = self._call("batch", cmds=';'.join(cmds))

        return [unescapearg(r) for r in rsp.split(';')]

    def _submitone(self, op, args):

        return self._call(op, **args)

    @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(urllib.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)

        return changegroupmod.cg1unpacker(f, 'UN')

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

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

        opts = {}

        bundlecaps = kwargs.get('bundlecaps')

        if bundlecaps is not None:

            kwargs['bundlecaps'] = sorted(bundlecaps)

        else:

            bundlecaps = () # kwargs could have it to None

        for key, value in kwargs.iteritems():

            if value is None:

                continue

            keytype = gboptsmap.get(key)

            if keytype is None:

                assert False, 'unexpected'

            elif keytype == 'nodes':

                value = encodelist(value)

            elif keytype in ('csv', 'scsv'):

                value = ','.join(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", **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, source):

        '''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 localrepository.addchangegroup()).

        When pushing a bundle20 stream, return a bundle20 stream.'''

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

            heads = encodelist(['hashed',

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