#!/usr/bin/env python3
#
# A tool to archive all of https://cr.joyent.us CRs to a local "./archive" dir.
#
# Assumptions:
# - An entry in "~/.ssh/config" equivalent to the following:
#        Host cr
#            User trentm
#            Hostname cr.joyent.us
#            Port 29418
#            ControlMaster no
# - A key that works for `ssh cr gerrit --help`
#
# Usage:
#   ./bin/cr.joyent.us-archive ARCHIVE-DIR
#

import argparse
import codecs
import json
import os
from pprint import pprint
from subprocess import run, PIPE
import sys
import time


#---- globals/config

# Typically good default config values.
skipRepoRefetch = False
gerritUser = os.environ['USER']  # Set manually if running as root somewhere.
hackRepoSkips = []  # Array of repos to exclude for a manual run.
cloneTimeoutS = 600


# These CRs have very large patch sets. We skip all but the "currentPatchSet"
# for these CRs.
patchSetSkipCrnums = [
    '2179',
    '2815',
    '3529',
    '3638',
    '4265',
    '4956',
    '5213',
    '5332',
    '6930',
    '3659',
    '1862',
    '1575',
    '3000',
    '3929',
    '3928',
    '4750',
    '720'
]

# Further, these CRs have a currentPatchSet so large that we don't want
# to keep it. The threshold here is >50MB at which point GitHub starts
# suggesting using "Git Large File Storage" https://git-lfs.github.com.
currentPatchSetSkipCrnums = [
    '2815'
]


# Projects (repos) to skip because of cloning problems with the cr.joyent.us
# data.
repoSkips = [
    # This fails:
    #   $ git clone trentm@cr.joyent.us:joyent/postgres.git
    #   Cloning into 'postgres'...
    #   remote: Counting objects: 649247, done
    #   fatal: internal server error
    #   remote: internal server error
    #   fatal: early EOF
    #   fatal: index-pack failed
    'joyent/postgres',

    # This hangs:
    #   $ git clone trentm@cr.joyent.us:All-Projects.git
    #   Cloning into 'All-Projects'...
    'All-Projects'
]

#---- support stuff

def clip(s, length):
    if len(s) <= length:
        return s
    else:
        return s[:length - 3] + '...'

def gerritQuery(queryArgv):
    sshArgv = [
        'ssh', '-q', '-o', 'TCPKeepAlive=yes', '-o', 'StrictHostKeyChecking=no',
        '-o', 'UserKnownHostsFile=/dev/null', '-o', 'ConnectTimeout=10',
        'cr', 'gerrit', 'query',
    ] + queryArgv
    res = run(sshArgv, stdout=PIPE, stderr=PIPE, check=True,
        universal_newlines=True)
    return res.stdout

def gerritGetChangeJson(number):
    output = gerritQuery([
        '--format=JSON',
        '--files', '--comments', '--all-reviewers', '--all-approvals',
        '--current-patch-set', '--', 'change:' + number])
    return output.splitlines()[0]

def updateBareRepo(project, tmpDir):
    repoUrl = '{}@cr.joyent.us:{}.git'.format(gerritUser, project)
    repoDir = os.path.join(tmpDir, 'repos', project + '.git')
    if os.path.isdir(repoDir):
        if not skipRepoRefetch:
            print('    "git fetch" in "{}"'.format(repoDir))
            run(['git', 'fetch'], cwd=repoDir,
                check=True, stdout=PIPE, stderr=PIPE, timeout=5)
    else:
        wrkDir = os.path.dirname(repoDir)
        if not os.path.isdir(wrkDir):
            os.makedirs(wrkDir)
        print('    "git clone --mirror {}" in "{}"'.format(repoUrl, wrkDir))
        run(['git', 'clone', '--mirror', repoUrl], cwd=wrkDir,
            check=True, timeout=cloneTimeoutS)
    return repoDir

def gitFormatPatch(repoDir, revision):
    # Explicitly do *not* convert this to string. We don't know an appropriate
    # encoding, in general.
    print('    "git format-patch -1 --stdout {}" in "{}"'.format(revision, repoDir))
    res = run(['git', 'format-patch', '-1', '--stdout', revision], cwd=repoDir,
        check=True, stdout=PIPE, stderr=PIPE)
    return res.stdout

def gitFormatPatchNoBinary(repoDir, revision):
    # Explicitly do *not* convert this to string. We don't know an appropriate
    # encoding, in general.
    print('    "git format-patch -1 --no-binary --stdout {}" in "{}"'.format(revision, repoDir))
    res = run(['git', 'format-patch', '-1', '--no-binary', '--stdout', revision], cwd=repoDir,
        check=True, stdout=PIPE, stderr=PIPE)
    return res.stdout

def archiveChange(change, archiveDir, tmpDir):
    if change['project'] in hackRepoSkips:
        print('    HACK: skip project {} for now'.format(change['project']))
        return

    changeDir = os.path.join(archiveDir, change['number'])
    if not os.path.isdir(changeDir):
        os.mkdir(changeDir)

    output = gerritGetChangeJson(change['number'])
    changeInfo = json.loads(output)
    # Alternatively, load the existing one:
    #  changePath = os.path.join(changeDir, 'change.json')
    #  with codecs.open(changePath, 'r', 'utf8') as f:
    #      changeInfo = json.loads(f.read())
    #pprint(changeInfo)

    # Binary files in `git format-patch` output gets HUGE. Let's skip those
    # unless the CR is open, and hence a candidate for `git am` usage.
    if changeInfo['open']:
        gitFormatPatchFn = gitFormatPatch
    else:
        gitFormatPatchFn = gitFormatPatchNoBinary

    # Gather diffs for patchsets.
    if changeInfo['project'] in repoSkips:
        print('    warning: skip project "{}" repo clone, because it fails'
            .format(changeInfo['project']))
    elif len(changeInfo['patchSets']) == 0:
        print('    no patchSets for this CR')
    else:
        print('    clone {}.git in {}'.format(changeInfo['project'], tmpDir))
        repoDir = updateBareRepo(changeInfo['project'], tmpDir)
        if changeInfo['number'] in patchSetSkipCrnums:
            print('    warning: skip CR {} patchSets because they are large'
                .format(changeInfo['number']))
        else:
            for patchSet in changeInfo['patchSets']:
                #pprint(patchSet)
                diff = gitFormatPatchFn(repoDir, patchSet['revision'])
                diffPath = os.path.join(changeDir,
                    'patchSet-{}.diff'.format(patchSet['number']))
                print('    write', diffPath)
                with codecs.open(diffPath, 'wb') as f:
                    f.write(diff)
        if changeInfo['number'] in currentPatchSetSkipCrnums:
            print('    warning: skip CR {} currentPatchSet because it is very large'
                .format(changeInfo['number']))
        else:
            # The "currentPatchSet" is a dupe, but can be convenient to have a
            # well-known name for the latest one.
            diff = gitFormatPatchFn(repoDir, changeInfo['currentPatchSet']['revision'])
            diffPath = os.path.join(changeDir, 'currentPatchSet.diff')
            print('    write', diffPath)
            with codecs.open(diffPath, 'wb') as f:
                f.write(diff)

    # Write the change.json last because it is used as the marker for
    # this dump having completed.
    changePath = os.path.join(changeDir, 'change.json')
    print('    write', changePath)
    with codecs.open(changePath, 'w', 'utf8') as f:
        f.write(output)


#---- mainline

def main(argv):
    parser = argparse.ArgumentParser(
        description='Archive all cr.joyent.us CRs.')
    parser.add_argument('--no-regen-changes', '-C', dest='regenerateAllChanges',
        action='store_false', default=True,
        help='Do not regenerate archive/all-changes.txt. This can be useful ' +
            'because regenerating this file can take a while and is flaky.')
    parser.add_argument('archiveDir', metavar="ARCHIVE-DIR")
    args = parser.parse_args()

    archiveDir = args.archiveDir
    if not os.path.isdir(archiveDir):
        os.makedirs(archiveDir)

    allChangesJson = os.path.join(archiveDir, 'all-changes.json')
    tmpDir = '/var/tmp/cr.joyent.us-archive'

    if args.regenerateAllChanges:
        # (Re-)generate archive/all-changes.json.
        print("Listing all changes")
        start = 0
        limit = 500
        moreChanges = True
        allChanges = []
        while moreChanges:
            print('  get one page: start={:4d} limit={}'.format(start, limit))
            output = gerritQuery([
                '--format=JSON',
                '--start={}'.format(start),
                '--',
                'age:0s',
                'limit:{}'.format(limit)
            ])
            for line in output.splitlines():
                record = json.loads(line)
                if record.get('type') == 'stats':
                    moreChanges = record['moreChanges']
                    start += limit
                else:
                    allChanges.append(record)

        print('  write', allChangesJson)
        with codecs.open(allChangesJson, 'w', 'utf8') as f:
            for change in allChanges:
                f.write(json.dumps(change) + '\n')
    else:
        # Load archive/all-changes.json from disk. This is typically *not*
        # done.
        allChanges = []
        with codecs.open(allChangesJson, 'r', 'utf8') as f:
            for line in f.readlines():
                allChanges.append(json.loads(line))

    # Update each change
    print('Update all changes')
    for i, change in enumerate(allChanges):
        print(
            '  [{:4d}/{}] https://cr.joyent.us/{} ... '.format(
                i + 1,
                len(allChanges),
                change['number']
            ),
            end=''
        )
        changeJson = os.path.join(archiveDir, change['number'], 'change.json')
        try:
            with codecs.open(changeJson, 'r', 'utf8') as f:
                archivedChange = json.load(f)
        except FileNotFoundError:
            archivedChange = None
        if archivedChange and archivedChange['lastUpdated'] >= change['lastUpdated']:
            print('up to date')
        else:
            print('need to update')
            archiveChange(change, archiveDir, tmpDir)


if __name__ == '__main__':
    sys.exit(main(sys.argv))
