#!/usr/bin/env python3

# 1) Create a base backup of the database on the template system.
# 2) Create a tar file of the paths given in the config
# 3) Transfer the tar file to the local system
# Configuration can be set in ``config/template_conf.py`` or
# ``config/source_conf.py``.
# Which configuration is to be used is given by the positional argument.


import subprocess as sp
import os
import argparse

from perfact.generic import load_config

config = None  # Will later be read


def remote_cmd(cmd):
    return ['ssh', config['system'], cmd]


def fetch(src, tgt):
    sp.run(
        ['/usr/bin/rsync', '--compress', '--progress', '--partial',
         '{system}:{src}'.format(src=src, **config),
         tgt],
        check=True,
    )


def cleanup(archive):
    sp.run(remote_cmd(f"sudo rm -rf {archive}"))


def fetch_system():
    print("Creating system archive")
    os.makedirs(os.path.dirname(config['target']), exist_ok=True)
    try:
        proc = sp.run(remote_cmd(
            "sudo tar"
            "    -cp --numeric-owner --one-file-system --sort=name"
            "    -f {archive} {excludes} {folders}".format(**config)
        ))
        # Allow exit code 1, which means files vanished before they could be
        # packed
        assert proc.returncode in (0, 1), "Unable to create system archive"
        fetch(config['archive'], config['target'])
    finally:
        cleanup(config['archive'])


def fetch_db():
    print("Creating database archive")
    os.makedirs(os.path.dirname(config['target_db']), exist_ok=True)
    try:
        sp.run(remote_cmd(f"sudo mkdir {config['basebackup_target']}"),
               check=True)
        sp.run(
            remote_cmd("sudo bash"),
            input=(
                "set -e; "
                "mkdir -p {basebackup_target}; "
                "chown postgres {basebackup_target}; "
                "sudo -u postgres "
                "  pg_basebackup --progress --checkpoint=fast --format=plain "
                "    -X stream --pgdata={basebackup_target}"
            ).format(**config),
            universal_newlines=True,
            check=True
        )
        sp.run(
            remote_cmd(
                "sudo tar -cp --numeric-owner --one-file-system --sort=name"
                " -C {basebackup_target} -f {archive_db} .".format(**config),
            ),
            check=True,
        )
        cleanup(config['basebackup_target'])
        fetch(config['archive_db'], config['target_db'])
    finally:
        cleanup(config['basebackup_target'])
        cleanup(config['archive_db'])


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'config',
        type=str,
        help="Config file to use. Should be config/template_conf.py or"
        " config/source_conf.py",
    )
    parser.add_argument('--skip-system', action='store_true', default=False)
    parser.add_argument('--skip-db', action='store_true', default=False)
    args = parser.parse_args()

    config = load_config(args.config)

    config['excludes'] = ' '.join([
        '--exclude ' + exclude
        for exclude in config['archive_excludes']
    ])
    config['folders'] = ' '.join(config['archive_folders'])
    if not args.skip_system:
        fetch_system()
    if not args.skip_db:
        fetch_db()