#!/usr/bin/env python3

import argparse
import os
import subprocess


def get_service_datasets(root_dataset):
    zfs_list = subprocess.getoutput(
        f"zfs list -H -r {root_dataset} -o name,mountpoint"
    )
    zfs_list_lines = zfs_list.split('\n')
    zfs_list_lines_items = map(lambda l: l.split(), zfs_list_lines)
    return {
        os.path.basename(dataset): {
            "dataset": dataset,
            "mountpoint": mountpoint,
        } for dataset, mountpoint in zfs_list_lines_items if dataset != root_dataset
    }

def get_last_daily_snapshot(dataset):
    snapshots = subprocess.getoutput(
        f"zfs list -t snapshot -H -r {dataset} -o name -s creation"
    )
    daily_snapshots = filter(lambda s: s.endswith("_daily"), snapshots.split('\n'))
    last_daily_snapshot = list(daily_snapshots)[-1]
    assert '@' in last_daily_snapshot
    assert last_daily_snapshot.split('@')[0] == dataset
    return last_daily_snapshot

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Backup service data using restic")
    parser.add_argument("--root-dataset", type=str, required=True,
                        help="The root data root whose children are to be backed up")
    parser.add_argument("--bucket-endpoint", type=str, required=True,
                        help="S3 bucket endpoint for the backups")
    args = parser.parse_args()

    service_datasets = get_service_datasets(args.root_dataset)
    for service, properties in service_datasets.items():
        properties["snapshot"] = get_last_daily_snapshot(properties["dataset"])

    for service, properties in service_datasets.items():
        mountpoint = properties["mountpoint"]
        snapshot = properties["snapshot"]
        backup_path = os.path.normpath(os.path.join("/", "mnt", os.path.relpath(mountpoint, "/")))
        restic_cmd_base = "restic " \
            f"--repo s3:https://{args.bucket_endpoint}/the-nine-worlds---{service} " \
            "--option s3.storage-class=ONEZONE_IA"

        print(f"Backing up {service} : {snapshot}", flush=True)

        subprocess.run(f"zfs clone -o mountpoint={backup_path} {snapshot} rpool/restic",
                       shell=True, check=True)
        try:
            subprocess.run(f"{restic_cmd_base} snapshots || {restic_cmd_base} init",
                           shell=True, check=True)
            subprocess.run(f"{restic_cmd_base} backup .", shell=True, cwd=backup_path, check=True)
            subprocess.run(f"{restic_cmd_base} forget --prune --keep-daily 90 --keep-monthly 12",
                           shell=True, check=True)
            subprocess.run(f"{restic_cmd_base} check", shell=True, check=True)

        finally:
            subprocess.run("zfs destroy rpool/restic", shell=True, check=True)