#!/usr/bin/env 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

SUPPORTED_ARCHS = {
    'amd64': 'qemu-system-x86_64',
    'i386': 'qemu-system-i386'
}

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(
        '--arch',
        action='store',
        help="Architecture to use (instead of attempting to auto-guess).",
    )
    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:
        print("Image does not exist", file=sys.stderr)
        sys.exit(1)

    if parsed_args.arch:
        try:
            qemu = SUPPORTED_ARCHS[parsed_args.arch]
        except KeyError as e:
            print(
                f"Unsupported architecture {parsed_args.arch}",
                file=sys.stderr,
            )
            sys.exit(1)
    else:
        # 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:
            print(
                f"Could not guess architecture, please use --arch",
                file=sys.stderr,
            )
            sys.exit(1)

    os.environ['DEBIAN_FRONTEND'] = 'noninteractive'
    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:~# ')
        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()
