# Copyright 2017 Facebook, Inc.

#

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

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

"""

    [infinitepushbackup]

    # Whether to enable automatic backups. If this option is True then a backup

    # process will be started after every mercurial command that modifies the

    # repo, for example, commit, amend, histedit, rebase etc.

    autobackup = False

    # path to the directory where pushback logs should be stored

    logdir = path/to/dir

    # Backup at most maxheadstobackup heads, other heads are ignored.

    # Negative number means backup everything.

    maxheadstobackup = -1

    # Nodes that should not be backed up. Ancestors of these nodes won't be

    # backed up either

    dontbackupnodes = []

    # Special option that may be used to trigger re-backuping. For example,

    # if there was a bug in infinitepush backups, then changing the value of

    # this option will force all clients to make a "clean" backup

    backupgeneration = 0

    # Hostname value to use. If not specified then socket.gethostname() will

    # be used

    hostname = ''

    # Enable reporting of infinitepush backup status as a summary at the end

    # of smartlog.

    enablestatus = False

    # Whether or not to save information about the latest successful backup.

    # This information includes the local revision number and unix timestamp

    # of the last time we successfully made a backup.

    savelatestbackupinfo = False

"""

from __future__ import absolute_import

import collections

import errno

import json

import os

import re

import socket

import stat

import subprocess

import time

from mercurial.node import (

    bin,

    hex,

    nullrev,

    short,

)

from mercurial.i18n import _

from mercurial import (

    bundle2,

    changegroup,

    commands,

    discovery,

    dispatch,

    encoding,

    error,

    extensions,

    hg,

    localrepo,

    lock as lockmod,

    phases,

    policy,

    registrar,

    scmutil,

    util,

)

from . import bundleparts

getscratchbookmarkspart = bundleparts.getscratchbookmarkspart

getscratchbranchparts = bundleparts.getscratchbranchparts

from hgext3rd import shareutil

osutil = policy.importmod(r'osutil')

cmdtable = {}

command = registrar.command(cmdtable)

revsetpredicate = registrar.revsetpredicate()

templatekeyword = registrar.templatekeyword()

backupbookmarktuple = collections.namedtuple('backupbookmarktuple',

                                 ['hostname', 'reporoot', 'localbookmark'])

class backupstate(object):

    def __init__(self):

        self.heads = set()

        self.localbookmarks = {}

    def empty(self):

        return not self.heads and not self.localbookmarks

class WrongPermissionsException(Exception):

    def __init__(self, logdir):

        self.logdir = logdir

restoreoptions = [

     ('', 'reporoot', '', 'root of the repo to restore'),

     ('', 'user', '', 'user who ran the backup'),

     ('', 'hostname', '', 'hostname of the repo to restore'),

]

_backuplockname = 'infinitepushbackup.lock'

def extsetup(ui):

    if ui.configbool('infinitepushbackup', 'autobackup', False):

        extensions.wrapfunction(dispatch, 'runcommand',

                                _autobackupruncommandwrapper)

        extensions.wrapfunction(localrepo.localrepository, 'transaction',

                                _transaction)

@command('pushbackup',

         [('', 'background', None, 'run backup in background')])

def backup(ui, repo, dest=None, **opts):

    """

    Pushes commits, bookmarks and heads to infinitepush.

    New non-extinct commits are saved since the last `hg pushbackup`

    or since 0 revision if this backup is the first.

    Local bookmarks are saved remotely as:

        infinitepush/backups/USERNAME/HOST/REPOROOT/bookmarks/LOCAL_BOOKMARK

    Local heads are saved remotely as:

        infinitepush/backups/USERNAME/HOST/REPOROOT/heads/HEAD_HASH

    """

    if opts.get('background'):

        _dobackgroundbackup(ui, repo, dest)

        return 0

    try:

        # Wait at most 30 seconds, because that's the average backup time

        timeout = 30

        srcrepo = shareutil.getsrcrepo(repo)

        with lockmod.lock(srcrepo.vfs, _backuplockname, timeout=timeout):

            return _dobackup(ui, repo, dest, **opts)

    except error.LockHeld as e:

        if e.errno == errno.ETIMEDOUT:

            ui.warn(_('timeout waiting on backup lock\n'))

            return 0

        else:

            raise

@command('pullbackup', restoreoptions)

def restore(ui, repo, dest=None, **opts):

    """

    Pulls commits from infinitepush that were previously saved with

    `hg pushbackup`.

    If user has only one backup for the `dest` repo then it will be restored.

    But user may have backed up many local repos that points to `dest` repo.

    These local repos may reside on different hosts or in different

    repo roots. It makes restore ambiguous; `--reporoot` and `--hostname`

    options are used to disambiguate.

    """

    other = _getremote(repo, ui, dest, **opts)

    sourcereporoot = opts.get('reporoot')

    sourcehostname = opts.get('hostname')

    namingmgr = BackupBookmarkNamingManager(ui, repo, opts.get('user'))

    allbackupstates = _downloadbackupstate(ui, other, sourcereporoot,

                                           sourcehostname, namingmgr)

    if len(allbackupstates) == 0:

        ui.warn(_('no backups found!'))

        return 1

    _checkbackupstates(allbackupstates)

    __, backupstate = allbackupstates.popitem()

    pullcmd, pullopts = _getcommandandoptions('^pull')

    # pull backuped heads and nodes that are pointed by bookmarks

    pullopts['rev'] = list(backupstate.heads |

                           set(backupstate.localbookmarks.values()))

    if dest:

        pullopts['source'] = dest

    result = pullcmd(ui, repo, **pullopts)

    with repo.wlock(), repo.lock(), repo.transaction('bookmark') as tr:

        changes = []

        for book, hexnode in backupstate.localbookmarks.iteritems():

            if hexnode in repo:

                changes.append((book, bin(hexnode)))

            else:

                ui.warn(_('%s not found, not creating %s bookmark') %

                        (hexnode, book))

        repo._bookmarks.applychanges(repo, tr, changes)

    # manually write local backup state and flag to not autobackup

    # just after we restored, which would be pointless

    _writelocalbackupstate(repo.vfs,

                           list(backupstate.heads),

                           backupstate.localbookmarks)

    repo.ignoreautobackup = True

    return result

@command('getavailablebackups',

    [('', 'user', '', _('username, defaults to current user')),

     ('', 'json', None, _('print available backups in json format'))])

def getavailablebackups(ui, repo, dest=None, **opts):

    other = _getremote(repo, ui, dest, **opts)

    sourcereporoot = opts.get('reporoot')

    sourcehostname = opts.get('hostname')

    namingmgr = BackupBookmarkNamingManager(ui, repo, opts.get('user'))

    allbackupstates = _downloadbackupstate(ui, other, sourcereporoot,

                                           sourcehostname, namingmgr)

    if opts.get('json'):

        jsondict = collections.defaultdict(list)

        for hostname, reporoot in allbackupstates.keys():

            jsondict[hostname].append(reporoot)

            # make sure the output is sorted. That's not an efficient way to

            # keep list sorted but we don't have that many backups.

            jsondict[hostname].sort()

        ui.write('%s\n' % json.dumps(jsondict))

    else:

        if not allbackupstates:

            ui.write(_('no backups available for %s\n') % namingmgr.username)

        ui.write(_('user %s has %d available backups:\n') %

                 (namingmgr.username, len(allbackupstates)))

        for hostname, reporoot in sorted(allbackupstates.keys()):

            ui.write(_('%s on %s\n') % (reporoot, hostname))

@command('debugcheckbackup',

         [('', 'all', None, _('check all backups that user have')),

         ] + restoreoptions)

def checkbackup(ui, repo, dest=None, **opts):

    """

    Checks that all the nodes that backup needs are available in bundlestore

    This command can check either specific backup (see restoreoptions) or all

    backups for the user

    """

    sourcereporoot = opts.get('reporoot')

    sourcehostname = opts.get('hostname')

    other = _getremote(repo, ui, dest, **opts)

    namingmgr = BackupBookmarkNamingManager(ui, repo, opts.get('user'))

    allbackupstates = _downloadbackupstate(ui, other, sourcereporoot,

                                           sourcehostname, namingmgr)

    if not opts.get('all'):

        _checkbackupstates(allbackupstates)

    ret = 0

    while allbackupstates:

        key, bkpstate = allbackupstates.popitem()

        ui.status(_('checking %s on %s\n') % (key[1], key[0]))

        if not _dobackupcheck(bkpstate, ui, repo, dest, **opts):

            ret = 255

    return ret

@command('debugwaitbackup', [('', 'timeout', '', 'timeout value')])

def waitbackup(ui, repo, timeout):

    try:

        if timeout:

            timeout = int(timeout)

        else:

            timeout = -1

    except ValueError:

        raise error.Abort('timeout should be integer')

    try:

        repo = shareutil.getsrcrepo(repo)

        with lockmod.lock(repo.vfs, _backuplockname, timeout=timeout):

            pass

    except error.LockHeld as e:

        if e.errno == errno.ETIMEDOUT:

            raise error.Abort(_('timeout while waiting for backup'))

        raise

@command('isbackedup',

     [('r', 'rev', [], _('show the specified revision or revset'), _('REV'))])

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

    """checks if commit was backed up to infinitepush

    If no revision are specified then it checks working copy parent

    """

    revs = opts.get('rev')

    if not revs:

        revs = ['.']

    bkpstate = _readlocalbackupstate(ui, repo)

    unfi = repo.unfiltered()

    backeduprevs = unfi.revs('draft() and ::%ls', bkpstate.heads)

    for r in scmutil.revrange(unfi, revs):

        ui.write(_(unfi[r].hex() + ' '))

        ui.write(_('backed up' if r in backeduprevs else 'not backed up'))

        ui.write(_('\n'))

@revsetpredicate('backedup')

def backedup(repo, subset, x):

    """Draft changesets that have been backed up by infinitepush"""

    unfi = repo.unfiltered()

    bkpstate = _readlocalbackupstate(repo.ui, repo)

    return subset & unfi.revs('draft() and ::%ls and not hidden()',

                              bkpstate.heads)

@revsetpredicate('notbackedup')

def notbackedup(repo, subset, x):

    """Changesets that have not yet been backed up by infinitepush"""

    bkpstate = _readlocalbackupstate(repo.ui, repo)

    bkpheads = set(bkpstate.heads)

    candidates = set(_backupheads(repo.ui, repo))

    notbackeduprevs = set()

    # Find all revisions that are ancestors of the expected backup heads,

    # stopping when we reach either a public commit or a known backup head.

    while candidates:

        candidate = candidates.pop()

        if candidate not in bkpheads:

            ctx = repo[candidate]

            rev = ctx.rev()

            if rev not in notbackeduprevs and ctx.phase() != phases.public:

                # This rev may not have been backed up.  Record it, and add its

                # parents as candidates.

                notbackeduprevs.add(rev)

                candidates.update([p.hex() for p in ctx.parents()])

    if notbackeduprevs:

        # Some revisions in this set may actually have been backed up by

        # virtue of being an ancestor of a different backup head, which may

        # have been hidden since the backup was made.  Find these and remove

        # them from the set.

        unfi = repo.unfiltered()

        candidates = bkpheads

        while candidates:

            candidate = candidates.pop()

            if candidate in unfi:

                ctx = unfi[candidate]

                if ctx.phase() != phases.public:

                    notbackeduprevs.discard(ctx.rev())

                    candidates.update([p.hex() for p in ctx.parents()])

    return subset & notbackeduprevs

@templatekeyword('backingup')

def backingup(repo, ctx, **args):

    """Whether infinitepush is currently backing up commits."""

    # If the backup lock exists then a backup should be in progress.

    srcrepo = shareutil.getsrcrepo(repo)

    return srcrepo.vfs.lexists(_backuplockname)

def smartlogsummary(ui, repo):

    if not ui.configbool('infinitepushbackup', 'enablestatus'):

        return

    # Don't output the summary if a backup is currently in progress.

    srcrepo = shareutil.getsrcrepo(repo)

    if srcrepo.vfs.lexists(_backuplockname):

        return

    unbackeduprevs = repo.revs('notbackedup()')

    # Count the number of changesets that haven't been backed up for 10 minutes.

    # If there is only one, also print out its hash.

    backuptime = time.time() - 10 * 60  # 10 minutes ago

    count = 0

    singleunbackeduprev = None

    for rev in unbackeduprevs:

        if repo[rev].date()[0] <= backuptime:

            singleunbackeduprev = rev

            count += 1

    if count > 0:

        if count > 1:

            ui.warn(_('note: %d changesets are not backed up.\n') % count)

        else:

            ui.warn(_('note: changeset %s is not backed up.\n') %

                    short(repo[singleunbackeduprev].node()))

        ui.warn(_('Run `hg pushbackup` to perform a backup.  If this fails,\n'

                  'please report to the Source Control @ FB group.\n'))

def _autobackupruncommandwrapper(orig, lui, repo, cmd, fullargs, *args):

    '''

    If this wrapper is enabled then auto backup is started after every command

    that modifies a repository.

    Since we don't want to start auto backup after read-only commands,

    then this wrapper checks if this command opened at least one transaction.

    If yes then background backup will be started.

    '''

    # For chg, do not wrap the "serve" runcommand call

    if 'CHGINTERNALMARK' in encoding.environ:

        return orig(lui, repo, cmd, fullargs, *args)

    try:

        return orig(lui, repo, cmd, fullargs, *args)

    finally:

        if getattr(repo, 'txnwasopened', False) \

                and not getattr(repo, 'ignoreautobackup', False):

            lui.debug("starting infinitepush autobackup in the background\n")

            _dobackgroundbackup(lui, repo)

def _transaction(orig, self, *args, **kwargs):

    ''' Wrapper that records if a transaction was opened.

    If a transaction was opened then we want to start background backup process.

    This hook records the fact that transaction was opened.

    '''

    self.txnwasopened = True

    return orig(self, *args, **kwargs)

def _backupheads(ui, repo):

    """Returns the set of heads that should be backed up in this repo."""

    maxheadstobackup = ui.configint('infinitepushbackup',

                                    'maxheadstobackup', -1)

    revset = 'heads(draft()) & not obsolete()'

    backupheads = [ctx.hex() for ctx in repo.set(revset)]

    if maxheadstobackup > 0:

        backupheads = backupheads[-maxheadstobackup:]

    elif maxheadstobackup == 0:

        backupheads = []

    return set(backupheads)

def _dobackup(ui, repo, dest, **opts):

    ui.status(_('starting backup %s\n') % time.strftime('%H:%M:%S %d %b %Y %Z'))

    start = time.time()

    # to handle multiple working copies correctly

    repo = shareutil.getsrcrepo(repo)

    currentbkpgenerationvalue = _readbackupgenerationfile(repo.vfs)

    newbkpgenerationvalue = ui.configint('infinitepushbackup',

                                         'backupgeneration', 0)

    if currentbkpgenerationvalue != newbkpgenerationvalue:

        # Unlinking local backup state will trigger re-backuping

        _deletebackupstate(repo)

        _writebackupgenerationfile(repo.vfs, newbkpgenerationvalue)

    bkpstate = _readlocalbackupstate(ui, repo)

    # this variable stores the local store info (tip numeric revision and date)

    # which we use to quickly tell if our backup is stale

    afterbackupinfo = _getlocalinfo(repo)

    # This variable will store what heads will be saved in backup state file

    # if backup finishes successfully

    afterbackupheads = _backupheads(ui, repo)

    other = _getremote(repo, ui, dest, **opts)

    outgoing, badhexnodes = _getrevstobackup(repo, ui, other,

                                             afterbackupheads - bkpstate.heads)

    # If remotefilelog extension is enabled then there can be nodes that we

    # can't backup. In this case let's remove them from afterbackupheads

    afterbackupheads.difference_update(badhexnodes)

    # As afterbackupheads this variable stores what heads will be saved in

    # backup state file if backup finishes successfully

    afterbackuplocalbooks = _getlocalbookmarks(repo)

    afterbackuplocalbooks = _filterbookmarks(

        afterbackuplocalbooks, repo, afterbackupheads)

    newheads = afterbackupheads - bkpstate.heads

    removedheads = bkpstate.heads - afterbackupheads

    newbookmarks = _dictdiff(afterbackuplocalbooks, bkpstate.localbookmarks)

    removedbookmarks = _dictdiff(bkpstate.localbookmarks, afterbackuplocalbooks)

    namingmgr = BackupBookmarkNamingManager(ui, repo)

    bookmarkstobackup = _getbookmarkstobackup(

        repo, newbookmarks, removedbookmarks,

        newheads, removedheads, namingmgr)

    # Special case if backup state is empty. Clean all backup bookmarks from the

    # server.

    if bkpstate.empty():

        bookmarkstobackup[namingmgr.getbackupheadprefix()] = ''

        bookmarkstobackup[namingmgr.getbackupbookmarkprefix()] = ''

    # Wrap deltaparent function to make sure that bundle takes less space

    # See _deltaparent comments for details

    extensions.wrapfunction(changegroup.cg2packer, 'deltaparent', _deltaparent)

    try:

        bundler = _createbundler(ui, repo, other)

        bundler.addparam("infinitepush", "True")

        backup = False

        if outgoing and outgoing.missing:

            backup = True

            parts = getscratchbranchparts(repo, other, outgoing,

                                          confignonforwardmove=False,

                                          ui=ui, bookmark=None,

                                          create=False)

            for part in parts:

                bundler.addpart(part)

        if bookmarkstobackup:

            backup = True

            bundler.addpart(getscratchbookmarkspart(other, bookmarkstobackup))

        if backup:

            _sendbundle(bundler, other)

            _writelocalbackupstate(repo.vfs, afterbackupheads,

                                   afterbackuplocalbooks)

            if ui.config('infinitepushbackup', 'savelatestbackupinfo'):

                _writelocalbackupinfo(repo.vfs, **afterbackupinfo)

        else:

            ui.status(_('nothing to backup\n'))

    finally:

        # cleanup ensures that all pipes are flushed

        cleanup = getattr(other, '_cleanup', None) or getattr(other, 'cleanup')

        try:

            cleanup()

        except Exception:

            ui.warn(_('remote connection cleanup failed\n'))

        ui.status(_('finished in %f seconds\n') % (time.time() - start))

        extensions.unwrapfunction(changegroup.cg2packer, 'deltaparent',

                                  _deltaparent)

    return 0

def _dobackgroundbackup(ui, repo, dest=None):

    background_cmd = ['hg', 'pushbackup']

    if dest:

        background_cmd.append(dest)

    logfile = None

    logdir = ui.config('infinitepushbackup', 'logdir')

    if logdir:

        # make newly created files and dirs non-writable

        oldumask = os.umask(0o022)

        try:

            try:

                username = util.shortuser(ui.username())