#!/usr/bin/env python3


import logging
import os
import random
import shlex
import shutil
import signal
import subprocess
import tempfile
import time


CLOUD_INIT_META_TEMPLATE = """\
# Amazon EC2 style metadata
local-hostname: {hostname}
"""

CLOUD_INIT_USER_TEMPLATE = """\
#cloud-config
ssh_authorized_keys:
- {pubkey}
"""


class QemuSystem:
    def __init__(self, name, base_image, disk_size, pubkey):
        self._name = name
        self._dirname = "."
        self._image = self._copy_image(base_image, disk_size)
        self._pubkey = pubkey
        self._memory = 512
        self._vcpus = 1
        self._username = None
        self._port = None
        self._pid = None

        fd, self._knownhosts = tempfile.mkstemp()
        os.close(fd)

    def _join(self, *names):
        return os.path.join(self._dirname, *names)

    def set_disk_size(self, size):
        self._disk_size = str(size * 1024 * 1024)

    def set_memory(self, memory):
        self._memory = memory

    def set_vcpus(self, vcpus):
        self._vcpus = vcpus

    def set_username(self, username):
        self._username = username

    def get_port(self):
        return self._port

    def _copy_image(self, base_image, size):
        image = self._join("qemu.img")
        logging.debug(f"QemuSystem: actual disk image: {image}")
        shutil.copy(base_image, image)
        if size is not None:
            subprocess.check_call(["qemu-img", "resize", "-q", image, str(size)])
        return image

    def _cloud_init_iso(self):
        iso = self._join("cloud-init.iso")

        config = self._join("cloud-init")
        os.mkdir(config)

        meta = os.path.join(config, "meta-data")
        write_file(meta, CLOUD_INIT_META_TEMPLATE.format(hostname=self._name))

        user = os.path.join(config, "user-data")
        write_file(user, CLOUD_INIT_USER_TEMPLATE.format(pubkey=self._pubkey))

        if os.path.exists(iso):
            os.remove(iso)
        subprocess.check_call(
            [
                "genisoimage",
                "-quiet",
                "-joliet",
                "-rock",
                "-output",
                iso,
                "-volid",
                "cidata",
                config,
            ]
        )

        return iso

    def start(self):
        iso = self._cloud_init_iso()

        self._port = random.randint(2000, 30000)
        pid_file = self._join("qemu.pid")
        out_file = self._join("qemu.out")
        err_file = self._join("qemu.err")

        MiB = 1024 ** 2
        memory = int(self._memory / MiB)

        argv = [
            "/usr/bin/daemonize",
            "-c",
            self._dirname,
            "-p",
            pid_file,
            "-e",
            err_file,
            "-o",
            out_file,
            "/usr/bin/qemu-system-x86_64",
            "-name",
            self._name,
            "-m",
            str(memory),
            "-cpu",
            "host",
            "-smp",
            f"cpus={self._vcpus}",
            "-drive",
            f"file={self._image},format=qcow2,cache=none,if=virtio",
            "-drive",
            f"file={iso},if=ide,media=cdrom",
            "-device",
            "virtio-net,netdev=user.0",
            "-netdev",
            f"user,id=user.0,hostfwd=tcp::{self._port}-:22",
            "-nographic",
            "-enable-kvm",
        ]

        subprocess.check_call(argv, stdout=None, stderr=None)
        self._pid = int(wait_for(lambda: got_pid(pid_file), 1))

        logging.debug("started qemu-system")
        logging.debug(f"  argv: {argv}")
        logging.debug(f"  pid: {self._pid}")

    def stop(self):
        if self._pid is not None:
            logging.debug(f"killing qemu-system process {self._pid}")
            os.kill(self._pid, signal.SIGTERM)

    def wait_for_ssh(self):
        def ssh_ok():
            output, exit = self.ssh(["true"])
            logging.debug(f"tried ssh, exit={exit}")
            if exit == 0:
                return True
            return None

        return wait_for(ssh_ok, 300, sleep=5)

    def ssh(self, argv):
        assert self._username is not None

        srcdir = globals()["srcdir"]

        ssh_opts = [
            "-i",
            os.path.join(srcdir, "ssh", "id"),
            "-p",
            str(self._port),
            "-l",
            self._username,
            f"-ouserknownhostsfile={self._knownhosts}",
            "-ostricthostkeychecking=accept-new",
            "-opasswordauthentication=no",
        ]

        real_argv = (
            ["ssh"] + ssh_opts + ["localhost", "--"] + [shlex.quote(x) for x in argv]
        )

        logging.debug(f"running on VM: {real_argv}")
        p = subprocess.Popen(
            real_argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
        )
        output, _ = p.communicate()
        logging.debug(f"exit code: {p.returncode}")
        logging.debug(f"output: {output}")
        return output, p.returncode


def read_file(filename):
    with open(filename, "r") as f:
        return f.read()


def write_file(filename, content):
    with open(filename, "w") as f:
        f.write(content)


def wait_for(func, timeout, sleep=0.1):
    start = time.time()
    while time.time() < start + timeout:
        val = func()
        if val is not None:
            return val
        time.sleep(sleep)


def got_pid(pid_file):
    if os.path.exists(pid_file):
        pid = read_file(pid_file)
        if pid:
            return pid.strip()