#!/usr/bin/env python

import os
import argparse
import sys
from datetime import (datetime,timezone)
from dateutil.tz import (tzlocal)
import subprocess
import json

from tempfile import TemporaryDirectory

import pathlib

import unshare
import signal


children = set()

def waitchildren():
    while children:
        waitpid, waitret = os.wait()
        children.remove(waitpid)


def borg_lv(lv, target, archive_prefix, dry_run, **args):
    vgp = lv.split('/')
    lvn = vgp[-1]
    snn = f'{lvn}_snap_{datetime.utcnow().strftime("%Y%m%dT%H%MZ")}'
    sn = "/".join([*vgp[:-1], snn])

    try:
        subprocess.run(['lvcreate', f'-l100%FREE', '-s', '-n', snn, lv], stdin=subprocess.DEVNULL, check=True)

        creation_time = None
        with subprocess.Popen(['lvs', '--reportformat=json', '-olv_time', sn], stdout=subprocess.PIPE) as proc:
            res = json.load(proc.stdout)
            creation_time = datetime.strptime(res['report'][0]['lv'][0]['lv_time'], '%Y-%m-%d %H:%M:%S %z').astimezone(timezone.utc)

        with TemporaryDirectory(prefix=f'{snn}_') as tmpdir:
            child = os.fork()
            if child == 0:
                unshare.unshare(unshare.CLONE_NEWNS)
                subprocess.run(['mount', '--make-rprivate', '/'], check=True)
                chroot = pathlib.Path(tmpdir) / 'chroot'
                upper = pathlib.Path(tmpdir) / 'upper'
                work = pathlib.Path(tmpdir) / 'work'
                for path in [chroot,upper,work]:
                    path.mkdir()
                subprocess.run(['mount', '-t', 'overlay', 'overlay', '-o', f'lowerdir=/,upperdir={upper},workdir={work}', chroot], check=True)
                bindMounts = ['nix', 'run', 'proc', 'dev', 'sys', pathlib.Path(os.path.expanduser('~')).relative_to('/')]
                if 'SSH_AUTH_SOCK' in os.environ:
                    bindMounts.append(pathlib.Path(os.environ['SSH_AUTH_SOCK']).parent.relative_to('/'))
                if 'BORG_CACHE_DIR' in os.environ:
                    bindMounts.append(pathlib.Path(os.environ['BORG_CACHE_DIR']).parent.relative_to('/'))
                for bindMount in bindMounts:
                    (chroot / bindMount).mkdir(parents=True,exist_ok=True)
                    subprocess.run(['mount', '--bind', pathlib.Path('/') / bindMount, chroot / bindMount], check=True)
                os.chroot(chroot)
                os.chdir('/')
                dir = pathlib.Path('/borg')
                dir.mkdir(parents=True,exist_ok=True)
                subprocess.run(['mount', f'/dev/{sn}', dir], check=True)

                archive_lvn=lvn
                if "/".join(vgp[:-1]) != archive_prefix.split('.')[-2]:
                    archive_lvn=f'{"_".join(vgp[:-1])}_{archive_lvn}'
                archive=f'{target}::{archive_prefix}{archive_lvn}-{creation_time.strftime("%Y-%m-%dT%H:%M:%S")}'

                env = os.environ.copy()
                create_args = ['borg',
                               'create',
                               '--lock-wait=600',
                               '--one-file-system',
                               '--compression=auto,zstd,10',
                               '--chunker-params=10,23,16,4095',
                               f'--timestamp={creation_time.strftime("%Y-%m-%dT%H:%M:%S")}',
                               '--show-rc',
                               # '--progress',
                               '--list',
                               '--filter=AMEi-?',
                               '--stats',
                               '--patterns-from=.backup-vidhar',
                               '--exclude-caches',
                               '--keep-exclude-tags'
                               ]
                env['BORG_FILES_CACHE_SUFFIX'] = "_".join(vgp)
                if dry_run:
                    create_args += ['--dry-run']
                create_args += [archive]
                subprocess.run(create_args, cwd=dir, check=True, env=env)

                os._exit(0)
            else:
                children.add(child)
                waitchildren()
    finally:
        waitchildren()
        subprocess.run(['lvremove', '--force', sn], stdin=subprocess.DEVNULL, check=True)

def sigterm(signum, frame):
    raise SystemExit(128 + signum)

def main():
    signal.signal(signal.SIGTERM, sigterm)
    
    parser = argparse.ArgumentParser()
    parser.add_argument('lv', metavar='LV')
    parser.add_argument('--target', metavar='REPO', default='borg.vidhar:.')
    parser.add_argument('--archive-prefix', metavar='REPO', default='yggdrasil.niflheim.ymir.')
    parser.add_argument('--dry-run', action='store_true')
    args = parser.parse_args()

    return borg_lv(**vars(args))

if __name__ == "__main__":
    sys.exit(main())