#! /usr/bin/python3
# SPDX-License-Identifier: BSD-3-Clause
# Author: Christian Kastner <ckk@kvr.at>


import argparse
import datetime
import os
import subprocess
import sys

import pexpect


IMAGEDIR = os.environ.get('IMAGEDIR', '/var/cache/qemu-sbuild')


def make_snapshot(image):
    iso_stamp = datetime.datetime.now().strftime('%Y-%m-%d_%H%M%S')
    run = subprocess.run(['qemu-img', 'snapshot', '-l', image], capture_output=True)
    tags = [t.split()[1].decode('utf-8') for t in run.stdout.splitlines()[2:]]

    if iso_stamp in tags:
        print(f'Error: snapshot for {iso_stamp} already exists.', file=sys.stderr)
        return False

    run = subprocess.run(['qemu-img', 'snapshot', '-c', iso_stamp, image])
    return True if run.returncode == 0 else False


def main():
    # init options
    parser = argparse.ArgumentParser(
        description="sbuild-update analog for QEMU images."
    )
    parser.add_argument('--snapshot', action='store_true',
        help="Create a snapshot of the image before updating it. Useful for "
             "reproducibility purposes."
    )
    parser.add_argument('--noexec', action='store_true',
        help="Don't actually do anything. Just print the command string that "
             "would be executed, and then exit.")
    parser.add_argument('image', action='store',
        help="Image path. Will be tried first as-is, and if no image exists "
             "at that location, then $IMAGEDIR/<image> is tried."
    )

    parsed_args = parser.parse_args()

    if os.path.exists(parsed_args.image):
        image = parsed_args.image
    elif os.path.exists(os.path.join(IMAGEDIR, parsed_args.image)):
        image = os.path.join(IMAGEDIR, parsed_args.image)
    else:
        raise ValueError('Image does not exist')

    # This assumes that images are named foo-bar-ARCH.img
    components = os.path.basename(parsed_args.image)[:-4].split('-')
    if 'amd64' in components:
        qemu = 'qemu-system-x86_64'
    elif 'i386' in components:
        qemu = 'qemu-system-i386'
    else:
        raise ValueError('Unsupported architecture')

    args = [
            qemu,
            '-enable-kvm',
            '-object', 'rng-random,filename=/dev/urandom,id=rng0',
            '-device', 'virtio-rng-pci,rng=rng0,id=rng-device0',
            '-m',      '2048',
            '-nographic',
        ]
    args.append(image)

    print(' '.join(str(a) for a in args))
    if not parsed_args.noexec:
        if parsed_args.snapshot and not make_snapshot(image):
            return
        child = pexpect.spawn(' '.join(args), encoding='utf-8')
        child.timeout = 240
        child.expect('host login: ')
        child.sendline('root')
        child.logfile = sys.stdout
        child.expect('root@host:~# ')
        child.sendline('apt-get --quiet update')
        child.expect('root@host:~# ')
        # dist-upgrade can still open dialogue boxes that are impossible to
        # cleak away. Perhaps each sendline should be followed up by an [Enter]
        child.sendline('apt-get --quiet --assume-yes upgrade')
        child.expect('root@host:~# ')
        child.sendline('apt-get --quiet --assume-yes clean')
        child.expect('root@host:~# ')
        child.sendline('apt-get --quiet --assume-yes autoremove')
        child.expect('root@host:~# ')
        child.sendline('sync')
        child.expect('root@host:~# ')
        # Don't recall what issue this solves, but it solves it
        child.sendline('sleep 1')
        child.expect('root@host:~# ')
        child.sendline('shutdown -h now')


if __name__ == '__main__':
    main()
