# Code dedicated to the discovery of obsolescence marker "over the wire"

#

# Copyright 2017 Pierre-Yves David <pierre-yves.david@ens-lyon.org>

#

# 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

try:

    import StringIO as io

    StringIO = io.StringIO

except ImportError:

    import io

    StringIO = io.StringIO

import collections

import hashlib

import heapq

import math

import struct

from mercurial import (

    bundle2,

    cmdutil,

    commands,

    dagutil,

    error,

    exchange,

    extensions,

    localrepo,

    node,

    obsolete,

    scmutil,

    setdiscovery,

    util,

    wireproto,

)

from mercurial.hgweb import hgweb_mod

from mercurial.i18n import _

from . import (

    exthelper,

    utility,

)

_pack = struct.pack

_unpack = struct.unpack

eh = exthelper.exthelper()

obsexcmsg = utility.obsexcmsg

##########################################

###  trigger discovery during exchange ###

##########################################

@eh.wrapfunction(exchange, '_pushdiscoveryobsmarkers')

def _pushdiscoveryobsmarkers(orig, pushop):

    if (obsolete.isenabled(pushop.repo, obsolete.exchangeopt)

        and pushop.repo.obsstore

        and 'obsolete' in pushop.remote.listkeys('namespaces')):

        repo = pushop.repo

        obsexcmsg(repo.ui, "computing relevant nodes\n")

        revs = list(repo.revs('::%ln', pushop.futureheads))

        unfi = repo.unfiltered()

        cl = unfi.changelog

        if not pushop.remote.capable('_evoext_obshash_0'):

            # do not trust core yet

            # return orig(pushop)

            nodes = [cl.node(r) for r in revs]

            if nodes:

                obsexcmsg(repo.ui, "computing markers relevant to %i nodes\n"

                                   % len(nodes))

                pushop.outobsmarkers = repo.obsstore.relevantmarkers(nodes)

            else:

                obsexcmsg(repo.ui, "markers already in sync\n")

                pushop.outobsmarkers = []

                pushop.outobsmarkers = repo.obsstore.relevantmarkers(nodes)

            return

        common = []

        missing = None

        obsexcmsg(repo.ui, "looking for common markers in %i nodes\n"

                           % len(revs))

        commonrevs = list(unfi.revs('::%ln', pushop.outgoing.commonheads))

        if _canobshashrange(repo, pushop.remote):

            missing = findmissingrange(pushop.ui, unfi, pushop.remote,

                                       commonrevs)

        else:

            common = findcommonobsmarkers(pushop.ui, unfi, pushop.remote,

                                          commonrevs)

        if missing is None:

            revs = list(unfi.revs('%ld - (::%ln)', revs, common))

            nodes = [cl.node(r) for r in revs]

        else:

            revs = list(repo.revs('only(%ln, %ln)', pushop.futureheads,

                        pushop.outgoing.commonheads))

            nodes = [cl.node(r) for r in revs]

            nodes += missing

        if nodes:

            obsexcmsg(repo.ui, "computing markers relevant to %i nodes\n"

                               % len(nodes))

            pushop.outobsmarkers = repo.obsstore.relevantmarkers(nodes)

        else:

            obsexcmsg(repo.ui, "markers already in sync\n")

            pushop.outobsmarkers = []

@eh.extsetup

def _installobsmarkersdiscovery(ui):

    olddisco = exchange.pushdiscoverymapping['obsmarker']

    def newdisco(pushop):

        _pushdiscoveryobsmarkers(olddisco, pushop)

    exchange.pushdiscoverymapping['obsmarker'] = newdisco

def buildpullobsmarkersboundaries(pullop, bundle2=True):

    """small function returning the argument for pull markers call

    may to contains 'heads' and 'common'. skip the key for None.

    It is a separed function to play around with strategy for that."""

    repo = pullop.repo

    remote = pullop.remote

    unfi = repo.unfiltered()

    revs = unfi.revs('::(%ln - null)', pullop.common)

    boundaries = {'heads': pullop.pulledsubset}

    if not revs: # nothing common

        boundaries['common'] = [node.nullid]

        return boundaries

    if bundle2 and _canobshashrange(repo, remote):

        obsexcmsg(repo.ui, "looking for common markers in %i nodes\n"

                  % len(revs))

        boundaries['missing'] = findmissingrange(repo.ui, repo, pullop.remote,

                                                 revs)

    elif remote.capable('_evoext_obshash_0'):

        obsexcmsg(repo.ui, "looking for common markers in %i nodes\n"

                           % len(revs))

        boundaries['common'] = findcommonobsmarkers(repo.ui, repo, remote, revs)

    else:

        boundaries['common'] = [node.nullid]

    return boundaries

##################################

###  Code performing discovery ###

##################################

def _canobshashrange(local, remote):

    return (local.ui.configbool('experimental', 'obshashrange', False)

            and remote.capable('_evoext_obshashrange_1'))

def _obshashrange_capabilities(orig, repo, proto):

    """wrapper to advertise new capability"""

    caps = orig(repo, proto)

    enabled = repo.ui.configbool('experimental', 'obshashrange', False)

    if obsolete.isenabled(repo, obsolete.exchangeopt) and enabled:

        caps = caps.split()

        caps.append('_evoext_obshashrange_1')

        caps.sort()

        caps = ' '.join(caps)

    return caps

@eh.extsetup

def obshashrange_extsetup(ui):

    extensions.wrapfunction(wireproto, 'capabilities', _obshashrange_capabilities)

    # wrap command content

    oldcap, args = wireproto.commands['capabilities']

    def newcap(repo, proto):

        return _obshashrange_capabilities(oldcap, repo, proto)

    wireproto.commands['capabilities'] = (newcap, args)

def findcommonobsmarkers(ui, local, remote, probeset,

                         initialsamplesize=100,

                         fullsamplesize=200):

    # from discovery

    roundtrips = 0

    cl = local.changelog

    dag = dagutil.revlogdag(cl)

    missing = set()

    common = set()

    undecided = set(probeset)

    totalnb = len(undecided)

    ui.progress(_("comparing with other"), 0, total=totalnb)

    _takefullsample = setdiscovery._takefullsample

    if remote.capable('_evoext_obshash_1'):

        getremotehash = remote.evoext_obshash1

        localhash = _obsrelsethashtreefm1(local)

    else:

        getremotehash = remote.evoext_obshash

        localhash = _obsrelsethashtreefm0(local)

    while undecided:

        ui.note(_("sampling from both directions\n"))

        if len(undecided) < fullsamplesize:

            sample = set(undecided)

        else:

            sample = _takefullsample(dag, undecided, size=fullsamplesize)

        roundtrips += 1

        ui.progress(_("comparing with other"), totalnb - len(undecided),

                    total=totalnb)

        ui.debug("query %i; still undecided: %i, sample size is: %i\n"

                 % (roundtrips, len(undecided), len(sample)))

        # indices between sample and externalized version must match

        sample = list(sample)

        remotehash = getremotehash(dag.externalizeall(sample))

        yesno = [localhash[ix][1] == remotehash[si]

                 for si, ix in enumerate(sample)]

        commoninsample = set(n for i, n in enumerate(sample) if yesno[i])

        common.update(dag.ancestorset(commoninsample, common))

        missinginsample = [n for i, n in enumerate(sample) if not yesno[i]]

        missing.update(dag.descendantset(missinginsample, missing))

        undecided.difference_update(missing)

        undecided.difference_update(common)

    ui.progress(_("comparing with other"), None)

    result = dag.headsetofconnecteds(common)

    ui.debug("%d total queries\n" % roundtrips)

    if not result:

        return set([node.nullid])

    return dag.externalizeall(result)

def findmissingrange(ui, local, remote, probeset,

                     initialsamplesize=100,

                     fullsamplesize=200):

    missing = set()

    heads = local.revs('heads(%ld)', probeset)

    # size of slice ?

    heappop = heapq.heappop

    heappush = heapq.heappush

    heapify = heapq.heapify

    tested = set()

    sample = []

    samplesize = initialsamplesize

    def addentry(entry):

        if entry in tested:

            return False

        sample.append(entry)

        tested.add(entry)

        return True

    for h in heads:

        entry = _range(local, h, 0)

        addentry(entry)

    querycount = 0

    ui.progress(_("comparing obsmarker with other"), querycount)

    overflow = []

    while sample or overflow:

        if overflow:

            sample.extend(overflow)

            overflow = []

        if samplesize < len(sample):

            # too much sample already

            overflow = sample[samplesize:]

            sample = sample[:samplesize]

        elif len(sample) < samplesize:

            # we need more sample !

            needed = samplesize - len(sample)

            sliceme = []

            heapify(sliceme)

            for entry in sample:

                if 1 < len(entry):

                    heappush(sliceme, (-len(entry), entry))

            while sliceme and 0 < needed:

                _key, target = heappop(sliceme)

                for new in target.subranges():

                    # XXX we could record hierarchy to optimise drop

                    if addentry(entry):

                        if 1 < len(entry):

                            heappush(sliceme, (-len(entry), entry))

                        needed -= 1

                        if needed <= 0:

                            break

        # no longer the first interation

        samplesize = fullsamplesize

        nbsample = len(sample)

        maxsize = max([len(r) for r in sample])

        ui.debug("query %i; sample size is %i, largest range %i\n"

                 % (querycount, maxsize, nbsample))

        nbreplies = 0

        replies = list(_queryrange(ui, local, remote, sample))

        sample = []

        for entry, remotehash in replies:

            nbreplies += 1

            if remotehash == entry.obshash:

                continue

            elif 1 == len(entry):

            else:

                for new in entry.subranges():

                    addentry(new)

        assert nbsample == nbreplies

        querycount += 1

        ui.progress(_("comparing obsmarker with other"), querycount)

    ui.progress(_("comparing obsmarker with other"), None)

def _queryrange(ui, repo, remote, allentries):

    mapping = {}

    n = repo.changelog.node

    def gen():

        for entry in allentries:

            key = n(entry.head) + _pack('>I', entry.index)

            mapping[key] = entry

            yield key

    bundler = bundle2.bundle20(ui, bundle2.bundle2caps(remote))

    capsblob = bundle2.encodecaps(bundle2.getrepocaps(repo))

    bundler.newpart('replycaps', data=capsblob)

    bundler.newpart('_evoexp_obsrangehash_0', data=gen())

    stream = util.chunkbuffer(bundler.getchunks())

    try:

        reply = remote.unbundle(

            stream, ['force'], remote.url())

    except error.BundleValueError as exc:

        raise error.Abort(_('missing support for %s') % exc)

    try:

        op = bundle2.processbundle(repo, reply)

    except error.BundleValueError as exc:

        raise error.Abort(_('missing support for %s') % exc)

    except bundle2.AbortFromPart as exc:

        ui.status(_('remote: %s\n') % exc)

        if exc.hint is not None:

            ui.status(_('remote: %s\n') % ('(%s)' % exc.hint))

        raise error.Abort(_('push failed on remote'))

    for rep in op.records['_evoexp_obsrangehash_0']:

        yield mapping[rep['key']], rep['value']

@bundle2.parthandler('_evoexp_obsrangehash_0', ())

def _processqueryrange(op, inpart):

    assert op.reply is not None

    replies = []

    data = inpart.read(24)

    while data:

        n = data[:20]

        index = _unpack('>I', data[20:])[0]

        r = op.repo.changelog.rev(n)

        rhash = _range(op.repo, r, index).obshash

        replies.append(data + rhash)

        data = inpart.read(24)

    op.reply.newpart('reply:_evoexp_obsrangehash_0', data=iter(replies))

@bundle2.parthandler('reply:_evoexp_obsrangehash_0', ())

def _processqueryrangereply(op, inpart):

    data = inpart.read(44)

    while data:

        key = data[:24]

        rhash = data[24:]

        op.records.add('_evoexp_obsrangehash_0', {'key': key, 'value': rhash})

        data = inpart.read(44)