#!/usr/bin/python3
from shell import *

source_dir    = '/'
target_dir    = None
target_device = None

copy_files = False
fix_fstab  = False
fix_boot   = False

use = "Use: bootbak [-s <source_root_dir> (default: /)] [<target_device>] <target_root_dir>\n" \
      "             [-c (copy files)] [-f (fix fstab)] [-b (install/update boot loader)] [-a (do all three)]\n" \
      "             [-t (test mode)] [-v (verbose)]\n" \
      "     Note: If target_device is provided, it will be partitioned, formated, and mounted at target_root_dir.\n" \
      "\n\n" \
      "     Examples:\n" \
      "           bootbak /bak -c                             # Rsync / to /bak\n"\
      "           bootbak -s /bak /dev/sdb /mnt/restore -a    # Reformat /dev/sdb as our new boot drive and restore it from /bak\n"\
      "           bootbak -s /bak /mnt/restore -a             # Same, but starting from a pre-formatted and mounted drive.\n"

def parseflag(arg):
    global copy_files, fix_fstab, fix_boot, source_dir
    if arg == '-c':
        copy_files = True
    elif arg == '-f':
        fix_fstab = True
    elif arg == '-b':
        fix_boot = True
    elif arg == '-a':
        copy_files = fix_fstab = fix_boot = True
    elif arg == '-s':
        source_dir = nextarg()
        assert source_dir == '/' or not source_dir.endswith('/')
    else:
        return True

def parsearg(arg):
    global target_dir, target_device
    if arg.startswith('/dev/'):
        assert not target_device
        target_device = arg
    elif target_dir is None:
        assert arg == '/' or not arg.endswith('/')
        target_dir = arg
    else:
        return True

verbose, testmode = argparse(parseflag, parsearg, use)

if not target_dir:
    bail(use)

import admin

try:
    if not isdir(target_dir):
        bail("%s is not a directory or doesn't exist."%(target_dir,))

    #================================================================
    # What kind of source is it (primary, or backup)?
    #
    source_drive_name = file_body_stripped(reroot(source_dir, '/drive.name'), "'root' or 'root_backup'")
    if source_drive_name not in ('root', 'root_backup'):
        bail("%s expected to contain 'root' or 'root_backup' but doesn't."%(reroot(source_dir, '/drive.name'),))

    #================================================================
    # Partition and format the drive from scratch if requested:
    #
    if target_device:
        admin.partition_format_and_mount_drive(target_device, target_dir, (source_drive_name=='root_backup') and 'restore' or 'root_backup')

    #================================================================
    # Sanity check target drive, and determine whether we're in backup or restore mode:
    #
    if copy_files or fix_fstab or fix_boot:

        target_drive_name = file_body_stripped(reroot(target_dir, '/drive.name'), "'root_backup' or 'restore'")

        if target_drive_name == 'root_backup':
            if source_drive_name != 'root':
                bail("According to /drive.name this would copy %r to %r which is probably not what you want."%(source_drive_name, target_drive_name), test_ignore=True)
            mode = 'backup'
        elif target_drive_name == 'restore':
            if source_drive_name != 'root_backup':
                bail("According to /drive.name this would copy %r to %r which is probably not what you want."%(source_drive_name, target_drive_name), test_ignore=True)
            mode = 'restore'
        else:
            bail("%s expected to contain 'root_backup' or 'restore' but doesn't."%(reroot(target_dir, '/drive.name'),))

        echo("Assuming %s mode."%(mode,))

    if fix_fstab or fix_boot:
        for mountpoint in ['/', '/boot/efi']:
            if not find_device('mountpoint', reroot(target_dir, mountpoint)):
                bail("ERROR: %s is not mounted as its own partition."%(reroot(target_dir, mountpoint),))

    #================================================================
    # Restore the system files...
    #
    if copy_files:

        if target_device and not user_wants_to("copy files from %s to %s"%(source_dir, target_dir)):
            bail("Aborting...")

        rsync(source_dir, target_dir, excludes=['/drive.name', '/nobak'])

    #================================================================
    # Fix /etc/fstab on target
    #
    if fix_fstab:

        admin.fix_fstab_UUIDs(target_dir)

    #================================================================
    # Install the boot loader
    #
    if fix_boot:

        if not user_wants_to("install/update boot loader onto %s (%s)"%(target_dir, device_name(find_device('mountpoint', target_dir)))):
            bail("Aborting...")

        admin.install_boot_loader(target_dir, nvram=(mode=='restore'))    # Backup mode doesn't change primary boot disk, but restore does.
 
        # install_boot_loader() exits when done, so we never reach here.
        assert False

except RunFail as e:
    bail(e)
except KeyboardInterrupt:
    bail("Aborted.")

