#!/usr/bin/python3

import os
import subprocess
import random
import string
import tempfile
import urllib.request
import math
import argparse

from invirt import database


def lvcreate(name, size):
    subprocess.run(['lvcreate', '-L', size, '-n', name, 'xenvg'],
                   check_output=True, encoding='utf-8', check=True)

def lvrename(dest, src):
    subprocess.run(['lvchange', '-an', f'xenvg/{src}'],
                   check_output=True, encoding='utf-8', check=True)

    subprocess.run(['lvrename', f'xenvg/{src}', f'xenvg/{dest}'],
                   check_output=True, encoding='utf-8', check=True)

    subprocess.run(['lvchange', '-ay', f'xenvg/{dest}'],
                   check_output=True, encoding='utf-8', check=True)

def lv_random(func, pattern, *args):
    """
    Run some LVM-related command with a random string in the LV name.

    func takes an LV name plus whatever's in *args and returns the
    return code of some LVM command, such as lvcreate or lvrename

    pattern must contain one '{}' pattern, which will be replaced
    by a 6-character random string.

    the script will attempt to re-run itself if the error code
    indicates that the destination already exists
    """

    # Keep trying until it works
    while True:
        letters = (random.choice(string.ascii_letters) for _ in range(6))
        rand_string = ''.join(letters)

        name = pattern.format(rand_string)

        try:
            func(name, *args)
        except subprocess.CalledProcessError as e:
            # 5 is the return code if the destination already exists
            if e.returncode != 5:
                raise
        else:
            return name

def fetch_image(cdrom):
    """
    Download a cdrom from a URI, shelling out to rsync if appropriate
    and otherwise trying to use urllib
    """

    full_uri = os.path.join(cdrom.mirror.uri_prefix, cdrom.uri_suffix)
    temp_file = tempfile.mkstemp()[1]
    print(f'Fetching image {cdrom.cdrom_id} from {full_uri} to {temp_file}')
    try:
        if full_uri.startswith('rsync://'):
            subprocess.run(['rsync', '--no-motd', '-tLP', full_uri, temp_file],
                           check_output=True, encoding='utf-8', check=True)
        else:
            # I'm not going to look for errors here, because I bet it'll
            # throw its own exceptions
            urllib.request.urlretrieve(full_uri, temp_file)
        return temp_file
    except:
        os.remove(temp_file)
        raise

def copy_file(src, dest):
    """
    Copy a file from one location to another using dd
    """

    subprocess.run(['dd', f'if={src}', f'of={dest}', 'bs=1M'],
                   check_output=True, encoding='utf-8', check=True)

def update_image(cdrom):
    """
    Update a cdrom image by downloading the latest version,
    transferring it into an LV, moving the old LV out of the way and
    the new LV into place
    """

    if cdrom.mirror_id is None:
        return

    temp_file = fetch_image(cdrom)

    st_size = os.stat(temp_file).st_size

    assert st_size > 0, 'CD-ROM image size is 0'

    megabytes = math.ceil((float(st_size) / (1024 * 1024)))
    cdrom_size = f'{megabytes}M'

    try:
        new_lv = lv_random(lvcreate, f'image-new_{cdrom.cdrom_id}' + '_{}', cdrom_size)
        copy_file(temp_file, f'/dev/xenvg/{new_lv}')
        lv_random(lvrename, f'image-old_{cdrom.cdrom_id}' + '_{}', 'image_{cdrom.cdrom_id}')
        lv_random(lvrename, f'image_{cdrom.cdrom_id}', new_lv)
        reap_images()
    finally:
        os.remove(temp_file)

def reap_images():
    """
    Remove stale cdrom images that are no longer in use

    update_image doesn't attempt to remove the old image because it
    might still be in use. reap_images attempts to delete any LVs
    starting with 'image-old_', but ignores errors, in case they're
    still being used.
    """

    lvm_list = subprocess.run(['lvs', '-o', 'lv_name', '--noheadings'],
                              check_output=True, encoding='utf-8', check=True)

    for lv in (s.strip() for s in lvm_list.stdout.readlines()):
        if lv.startswith('image-old_'):
            subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
            subprocess.run(['lvchange', '-a', 'n', f'/dev/xenvg/{lv}'])
            subprocess.run(['lvchange', '-a', 'ey', f'/dev/xenvg/{lv}'])
            subprocess.run(['lvremove', '--force', f'/dev/xenvg/{lv}'])

def action_add_cdrom(args):
    attrs = {key: vars(args)[key] for key in ('cdrom_id', 'description', 'mirror_id', 'uri_suffix')}

    cdrom = database.CDROM(**attrs)
    database.session.add(cdrom)
    database.session.commit()

    update_image(cdrom)

def action_add_mirror(args):
    attrs = {key: vars(args)[key] for key in ('mirror_id', 'uri_prefix')}

    mirror = database.Mirror(**attrs)
    database.session.add(mirror)
    database.session.commit()

def action_update(args):
    if not args.names:
        images = database.CDROM.query().all()
    else:
        images = [database.CDROM.query().get(arg) for arg in args.names]

    for cdrom in images:
        if cdrom is not None:
            update_image(cdrom)

def action_reap(args):
    reap_images()

def main():
    database.connect()
    database.session.begin()

    parser = argparse.ArgumentParser(description='Perform actions on the CD-ROM images in the database')
    subparsers = parser.add_subparsers(help='Action to perform')

    add_parser = subparsers.add_parser('add', help='Add new image to database')
    add_parser.set_defaults(func=lambda args: add_parser.print_help())
    add_subparsers = add_parser.add_subparsers()

    add_cdrom_parser = add_subparsers.add_parser('cdrom')
    add_cdrom_parser.add_argument('cdrom_id')
    add_cdrom_parser.add_argument('description')
    add_cdrom_parser.add_argument('mirror_id')
    add_cdrom_parser.add_argument('uri_suffix')
    add_cdrom_parser.set_defaults(func=action_add_cdrom)

    add_mirror_parser = add_subparsers.add_parser('mirror')
    add_mirror_parser.add_argument('mirror_id')
    add_mirror_parser.add_argument('uri_prefix')
    add_mirror_parser.set_defaults(func=action_add_mirror)

    update_parser = subparsers.add_parser('update', help='Update images in database')
    update_parser.add_argument('names', nargs='*', metavar='name', help='Shortnames of images to update')
    update_parser.set_defaults(func=action_update)

    reap_parser = subparsers.add_parser('reap', help='Reap old images in database')
    reap_parser.set_defaults(func=action_reap)

    args = parser.parse_args()
    if 'func' in args:
        args.func(args)
    else:
        parser.print_help()


if __name__ == '__main__':
    main()