# Copyright 2016-present Facebook. All Rights Reserved.

#

# protocol: logic for a server providing fastannotate support

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

import contextlib

import os

from mercurial.i18n import _

from mercurial.pycompat import open

from mercurial import (

    error,

    extensions,

    hg,

    util,

    wireprotov1peer,

    wireprotov1server,

)

from mercurial.utils import (

    urlutil,

)

from . import context

# common

def _getmaster(ui):

    """get the mainbranch, and enforce it is set"""

    master = ui.config(b'fastannotate', b'mainbranch')

    if not master:

        raise error.Abort(

            _(

                b'fastannotate.mainbranch is required '

                b'for both the client and the server'

            )

        )

    return master

# server-side

def _capabilities(orig, repo, proto):

    result = orig(repo, proto)

    result.append(b'getannotate')

    return result

def _getannotate(repo, proto, path, lastnode):

    # output:

    #   FILE := vfspath + '\0' + str(size) + '\0' + content

    #   OUTPUT := '' | FILE + OUTPUT

    result = b''

    buildondemand = repo.ui.configbool(

        b'fastannotate', b'serverbuildondemand', True

    )

    with context.annotatecontext(repo, path) as actx:

        if buildondemand:

            # update before responding to the client

            master = _getmaster(repo.ui)

            try:

                if not actx.isuptodate(master):

                    actx.annotate(master, master)

            except Exception:

                # non-fast-forward move or corrupted. rebuild automically.

                actx.rebuild()

                try:

                    actx.annotate(master, master)

                except Exception:

                    actx.rebuild()  # delete files

            finally:

                # although the "with" context will also do a close/flush, we

                # need to do it early so we can send the correct respond to

                # client.

                actx.close()

        # send back the full content of revmap and linelog, in the future we

        # may want to do some rsync-like fancy updating.

        # the lastnode check is not necessary if the client and the server

        # agree where the main branch is.

        if actx.lastnode != lastnode:

            for p in [actx.revmappath, actx.linelogpath]:

                if not os.path.exists(p):

                    continue

                with open(p, b'rb') as f:

                    content = f.read()

                vfsbaselen = len(repo.vfs.base + b'/')

                relpath = p[vfsbaselen:]

                result += b'%s\0%d\0%s' % (relpath, len(content), content)

    return result

def _registerwireprotocommand():

    if b'getannotate' in wireprotov1server.commands:

        return

    wireprotov1server.wireprotocommand(b'getannotate', b'path lastnode')(

        _getannotate

    )

def serveruisetup(ui):

    _registerwireprotocommand()

    extensions.wrapfunction(wireprotov1server, '_capabilities', _capabilities)

# client-side

def _parseresponse(payload):

    result = {}

    i = 0

    l = len(payload) - 1

    state = 0  # 0: vfspath, 1: size

    vfspath = size = b''

    while i < l:

        ch = payload[i : i + 1]

        if ch == b'\0':

            if state == 1:

                result[vfspath] = payload[i + 1 : i + 1 + int(size)]

                i += int(size)

                state = 0

                vfspath = size = b''

            elif state == 0:

                state = 1

        else:

            if state == 1:

                size += ch

            elif state == 0:

                vfspath += ch

        i += 1

    return result

def peersetup(ui, peer):

    class fastannotatepeer(peer.__class__):

        @wireprotov1peer.batchable

        def getannotate(self, path, lastnode=None):

            if not self.capable(b'getannotate'):

                ui.warn(_(b'remote peer cannot provide annotate cache\n'))

                return None, None

            else:

                args = {b'path': path, b'lastnode': lastnode or b''}

                return args, _parseresponse

    peer.__class__ = fastannotatepeer

@contextlib.contextmanager

def annotatepeer(repo):

    ui = repo.ui

    remotedest = ui.config(b'fastannotate', b'remotepath', b'default')

    remotepath = urlutil.get_unique_pull_path_obj(

        b'fastannotate',

        ui,

        remotedest,

    )

    peer = hg.peer(ui, {}, remotepath)

    try:

        yield peer

    finally:

        peer.close()

def clientfetch(repo, paths, lastnodemap=None, peer=None):

    """download annotate cache from the server for paths"""

    if not paths:

        return

    if peer is None:

        with annotatepeer(repo) as peer:

            return clientfetch(repo, paths, lastnodemap, peer)

    if lastnodemap is None:

        lastnodemap = {}

    ui = repo.ui

    results = []

    with peer.commandexecutor() as batcher:

        ui.debug(b'fastannotate: requesting %d files\n' % len(paths))

        for p in paths:

            results.append(

                batcher.callcommand(

                    b'getannotate',

                    {b'path': p, b'lastnode': lastnodemap.get(p)},

                )

            )

        for result in results:

            r = result.result()

            # TODO: pconvert these paths on the server?

            r = {util.pconvert(p): v for p, v in r.items()}

            for path in sorted(r):

                # ignore malicious paths

                if not path.startswith(b'fastannotate/') or b'/../' in (

                    path + b'/'

                ):

                    ui.debug(

                        b'fastannotate: ignored malicious path %s\n' % path

                    )

                    continue

                content = r[path]

                if ui.debugflag:

                    ui.debug(

                        b'fastannotate: writing %d bytes to %s\n'

                        % (len(content), path)

                    )

                repo.vfs.makedirs(os.path.dirname(path))

                with repo.vfs(path, b'wb') as f:

                    f.write(content)

def _filterfetchpaths(repo, paths):

    """return a subset of paths whose history is long and need to fetch linelog

    from the server. works with remotefilelog and non-remotefilelog repos.

    """

    threshold = repo.ui.configint(b'fastannotate', b'clientfetchthreshold', 10)

    if threshold <= 0:

        return paths

    result = []

    for path in paths:

        try:

            if len(repo.file(path)) >= threshold:

                result.append(path)

        except Exception:  # file not found etc.

            result.append(path)

    return result

def localreposetup(ui, repo):

    class fastannotaterepo(repo.__class__):

        def prefetchfastannotate(self, paths, peer=None):

            master = _getmaster(self.ui)

            needupdatepaths = []

            lastnodemap = {}

            try:

                for path in _filterfetchpaths(self, paths):

                    with context.annotatecontext(self, path) as actx:

                        if not actx.isuptodate(master, strict=False):

                            needupdatepaths.append(path)

                            lastnodemap[path] = actx.lastnode

                if needupdatepaths:

                    clientfetch(self, needupdatepaths, lastnodemap, peer)

            except Exception as ex:

                # could be directory not writable or so, not fatal

                self.ui.debug(b'fastannotate: prefetch failed: %r\n' % ex)

    repo.__class__ = fastannotaterepo

def clientreposetup(ui, repo):

    _registerwireprotocommand()

    if repo.local():

        localreposetup(ui, repo)

    # TODO: this mutates global state, but only if at least one repo

    # has the extension enabled. This is probably bad for hgweb.

    if peersetup not in hg.wirepeersetupfuncs:

        hg.wirepeersetupfuncs.append(peersetup)