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

import sqlite3

import struct

import time

import weakref

from mercurial import (

    dagutil,

    error,

    exchange,

    extensions,

    localrepo,

    node,

    obsolete,

    pycompat,

    scmutil,

    setdiscovery,

    util,

    wireproto,

)

from mercurial.hgweb import hgweb_mod

from mercurial.i18n import _

from . import (

    exthelper,

    obscache,

    utility,

    stablerange,

)

# prior to hg-4.2 there are not util.timer

if util.safehasattr(util, 'timer'):

    timer = util.timer

elif util.safehasattr(time, "perf_counter"):

    timer = time.perf_counter

elif getattr(pycompat, 'osname', os.name) == 'nt':

    timer = time.clock

else:

    timer = time.time

_pack = struct.pack

_unpack = struct.unpack

_calcsize = struct.calcsize

eh = exthelper.exthelper()

eh.merge(stablerange.eh)

obsexcmsg = utility.obsexcmsg

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

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

    starttime = timer()

    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)

    local.obsstore.rangeobshashcache.update(local)

    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)

    duration = timer() - starttime

    logmsg = ('obsdiscovery, %d/%d mismatch'

              ' - %d obshashrange queries in %.4f seconds\n')

    logmsg %= (len(missing), len(probeset), querycount, duration)

    ui.log('evoext-obsdiscovery', logmsg)

    ui.debug(logmsg)

    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)

    repo.obsstore.rangeobshashcache.update(repo)

    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

    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)

            pieces.append(mbin)

        pieces.sort()

    else:

        for subrange in repo.stablerange.subranges(repo, rangeid):

            obshash = _obshashrange(repo, subrange)

            if obshash != nullid:

                pieces.append(obshash)

    sha = hashlib.sha1()

    # note: if there is only one subrange with actual data, we'll just

    # reuse the same hash.

    if not pieces:

        obshash = node.nullid

    elif len(pieces) != 1 or obshash is None:

        sha = hashlib.sha1()

        for p in pieces:

            sha.update(p)

        obshash = sha.digest()

    cache[rangeid] = obshash

    return obshash

### sqlite caching

_sqliteschema = [

    """CREATE TABLE meta(schemaversion INTEGER NOT NULL,

                         tiprev        INTEGER NOT NULL,

                         tipnode       BLOB    NOT NULL,

                         nbobsmarker   INTEGER NOT NULL,

                         obssize       BLOB    NOT NULL,

                         obskey        BLOB    NOT NULL

                        );""",

    """CREATE TABLE obshashrange(rev     INTEGER NOT NULL,

                                 idx     INTEGER NOT NULL,

                                 obshash BLOB    NOT NULL,

                                 PRIMARY KEY(rev, idx));""",

    "CREATE INDEX range_index ON obshashrange(rev, idx);",

]

_queryexist = "SELECT name FROM sqlite_master WHERE type='table' AND name='meta';"

_clearmeta = """DELETE FROM meta;"""

_newmeta = """INSERT INTO meta (schemaversion, tiprev, tipnode, nbobsmarker, obssize, obskey)

            VALUES (?,?,?,?,?,?);"""

_updateobshash = "INSERT INTO obshashrange(rev, idx, obshash) VALUES (?,?,?);"

_querymeta = "SELECT schemaversion, tiprev, tipnode, nbobsmarker, obssize, obskey FROM meta;"

_queryobshash = "SELECT obshash FROM obshashrange WHERE (rev = ? AND idx = ?);"

_reset = "DELETE FROM obshashrange;"

class _obshashcache(obscache.dualsourcecache):

    _schemaversion = 2

    _cachename = 'evo-ext-obshashrange' # used for error message

    def __init__(self, repo):

        super(_obshashcache, self).__init__()

        self._vfs = repo.vfs

        self._path = repo.vfs.join('cache/evoext_obshashrange_v1.sqlite')

        self._new = set()

        self._valid = True

        self._repo = weakref.ref(repo.unfiltered())

        # cache status

        self._ondiskcachekey = None

        self._data = {}

    def clear(self, reset=False):

        super(_obshashcache, self).clear(reset=reset)

        self._data.clear()

        self._new.clear()

        if reset:

            self._valid = False

        if '_con' in vars(self):

            del self._con

    def get(self, rangeid):

        # revision should be covered by the tiprev

        #

        # XXX there are issue with cache warming, we hack around it for now