#!/usr/bin/python3

import os
import argparse
import subprocess as sp

archive = '/vol/backup/dist-upgrade-state.tar.gz'


def run(*cmd, **kw):
    "Wrapper for sp.run that defaults to check=True"
    args = {'check': True}
    args.update(kw)
    return sp.run(cmd, **args)


def pack():
    "Pack local state of installed packages"
    assert os.geteuid() == 0, "Must be root"
    run('tar', 'czf', archive, '/var/lib/apt/lists',
        '/var/cache/apt/archives')


def fetch():
    "Fetch archive from source-system"
    run('ssh', 'source-system',
        'sudo rsync -aP --remove-source-files'
        ' {archive} synctarget:{archive}'.format(archive=archive)
        )


def apply():
    "Unpack archive and apply dist-upgrade"
    print("Cleaning local package cache")
    run('apt-get', 'clean')
    assert os.geteuid() == 0, "Must be root"
    run('tar', 'xvf', archive, '-C', '/')
    print("Install postgres updates first (dry-run)")
    run(
        'apt-get install --dry-run --only-upgrade "postgresql*"'
        '| grep -v -E "Skipping"', shell=True
    )
    input("Hit enter to continue with actual run, Ctrl+C to abort")
    # Generate the package list and output
    run(
        'apt-get install --assume-no --only-upgrade "postgresql*"'
        '| grep -v -E "Skipping"', shell=True
    )
    user_input = input('Do you want to continue (otherwise no updates at all) [y/N]? ')
    # Start the installation regardless of user input
    # after manually verifying the desired action
    if user_input == 'y':
        run(
            'apt-get install --assume-yes --only-upgrade "postgresql*"'
            '| grep -v -E "Skipping"', shell=True
        )
    else:
        quit()
    print("\nRunning dist-upgrade for remaining packages (dry-run)")
    run('apt-get', 'dist-upgrade', '--dry-run')
    input("Hit enter to continue with actual run, Ctrl+C to abort")
    run('apt-get', 'dist-upgrade')
    run('apt-get', 'autoremove', '--purge')
    input(
        '\n'
        + "Cleaning up {archive} next.\n".format(archive=archive)
        + "If you still need it for the next system in the chain, hit Ctrl+C."
    )
    os.remove(archive)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        "Helper to upgrade Debian-based systems to the same state"
    )
    subs = parser.add_subparsers()
    for cmd in ['pack', 'fetch', 'apply']:
        runner = globals().get(cmd)  # Function defined above

        sub = subs.add_parser(cmd, help=runner.__doc__)
        if cmd in ['pack', 'fetch']:
            sub.add_argument('--force', action='store_true', default=False,
                             help='Overwrite existing file')
        sub.set_defaults(runner=runner)

    args = parser.parse_args()

    assert not os.path.exists(archive) or getattr(args, 'force', True), (
        "Archive exists. Add --force to overwrite"
    )
    if not hasattr(args, 'runner'):
        print('Use --help to show usage')
    else:
        try:
            args.runner()
        except KeyboardInterrupt:
            # Avoid printing traceback, just add newline
            print()
