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

# Status: Experiment in progress // open question

#

#   The final discovery algorithm and protocol will go into core when we'll be

#   happy with it.

#

#   Some of the code in this module is for compatiblity with older version

#   of evolve and will be eventually dropped.

from __future__ import absolute_import

try:

    import StringIO as io

    StringIO = io.StringIO

except ImportError:

    import io

    StringIO = io.StringIO

import hashlib

import heapq

import sqlite3

import struct

import weakref

from mercurial import (

    dagutil,

    error,

    exchange,

    extensions,

    localrepo,

    node,

    obsolete,

    scmutil,

    setdiscovery,

    util,

    wireproto,

)

from mercurial.hgweb import hgweb_mod

from mercurial.i18n import _

from . import (

    exthelper,

    obscache,

    utility,

    stablerange,

)

_pack = struct.pack

_unpack = struct.unpack

_calcsize = struct.calcsize

eh = exthelper.exthelper()

eh.merge(stablerange.eh)

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

    local.stablerange.warmup(local)

    rangelength = local.stablerange.rangelength

    subranges = local.stablerange.subranges

    # 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 = (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:

            ui.debug("query %i; add more sample (target %i, current %i)\n"

                     % (querycount, samplesize, len(sample)))

            # we need more sample !

            needed = samplesize - len(sample)

            sliceme = []

            heapify(sliceme)

            for entry in sample:

                if 1 < rangelength(local, entry):

                    heappush(sliceme, (-rangelength(local, entry), entry))

            while sliceme and 0 < needed:

                _key, target = heappop(sliceme)

                for new in subranges(local, target):

                    # XXX we could record hierarchy to optimise drop

                    if addentry(new):

                        if 1 < len(new):

                            heappush(sliceme, (-rangelength(local, new), new))

                        needed -= 1

                        if needed <= 0:

                            break

        # no longer the first interation

        samplesize = fullsamplesize

        nbsample = len(sample)

        maxsize = max([rangelength(local, r) for r in sample])

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

                 % (querycount, nbsample, maxsize))

        nbreplies = 0

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

        sample = []

        n = local.changelog.node

        for entry, remotehash in replies:

            nbreplies += 1

            if remotehash == _obshashrange(local, entry):

                continue

            elif 1 == rangelength(local, entry):

                missing.add(n(entry[0]))

            else:

                for new in subranges(local, entry):

                    addentry(new)

        assert nbsample == nbreplies

        querycount += 1

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

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

    local.obsstore.rangeobshashcache.save(local)

    return sorted(missing)

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

    #  question are asked with node

    n = repo.changelog.node

    noderanges = [(n(entry[0]), entry[1]) for entry in allentries]

    replies = remote.evoext_obshashrange_v0(noderanges)

    result = []

    for idx, entry in enumerate(allentries):

        result.append((entry, replies[idx]))

    return result

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

### Range Hash computation ###

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

@eh.command(

    'debugobshashrange',

    [

        ('', 'rev', [], 'display obshash for all (rev, 0) range in REVS'),

        ('', 'subranges', False, 'display all subranges'),

    ],

    _(''))

def debugobshashrange(ui, repo, **opts):

    """display the ::REVS set topologically sorted in a stable way

    """

    s = node.short

    revs = scmutil.revrange(repo, opts['rev'])

    # prewarm depth cache

    if revs:

        repo.stablerange.warmup(repo, max(revs))

    cl = repo.changelog

    rangelength = repo.stablerange.rangelength

    depthrev = repo.stablerange.depthrev

    if opts['subranges']:

        ranges = stablerange.subrangesclosure(repo, revs)

    else:

        ranges = [(r, 0) for r in revs]

    headers = ('rev', 'node', 'index', 'size', 'depth', 'obshash')

    linetemplate = '%12d %12s %12d %12d %12d %12s\n'

    headertemplate = linetemplate.replace('d', 's')

    ui.status(headertemplate % headers)

    for r in ranges:

        d = (r[0],

             s(cl.node(r[0])),

             r[1],

             rangelength(repo, r),

             depthrev(repo, r[0]),

             node.short(_obshashrange(repo, r)))

        ui.status(linetemplate % d)

    repo.obsstore.rangeobshashcache.save(repo)

def _obshashrange(repo, rangeid):

    """return the obsolete hash associated to a range"""

    cache = repo.obsstore.rangeobshashcache

    cl = repo.changelog

    cache.update(repo)

    obshash = cache.get(rangeid)

    if obshash is not None:

        return obshash

    pieces = []

    nullid = node.nullid

    if repo.stablerange.rangelength(repo, rangeid) == 1:

        rangenode = cl.node(rangeid[0])

        tmarkers = repo.obsstore.relevantmarkers([rangenode])

        pieces = []

        for m in tmarkers:

            mbin = obsolete._fm1encodeonemarker(m)