#!/usr/bin/python3

import argparse
import os
import shutil
import subprocess
import sys
import tempfile
import yaml


KERNELS = {
    "amd64": "linux-image-amd64",
    "arm64": "linux-image-arm64",
    "armhf": "linux-image-armmp-lpae",
    "i386": "linux-image-686",
    "ppc64el": "linux-image-powerpc64le",
}

QEMUS = {
    "amd64": "qemu-system-x86_64",
}


class Config:
    def __init__(self):
        p = argparse.ArgumentParser()
        p.add_argument("--dump", action="store_true")
        p.add_argument("--vmdb", action="store")
        p.add_argument("--tarball-directory", action="store", default=".")
        p.add_argument("--debian-release", action="store")
        p.add_argument("--grub", action="store", choices=["bios", "uefi", "ieee1275"])
        p.add_argument("--mklabel", action="store", choices=["msdos", "gpt"])
        p.add_argument("--arch", action="store")
        p.add_argument("--boot", action="store_true")
        p.add_argument("--maybe-boot", action="store_true")
        p.add_argument("--verbose", action="store_true")
        args = p.parse_args()

        self.dump = args.dump
        self.tarball_directory = args.tarball_directory
        self.vmdb_filename = args.vmdb
        self.vmdb = yaml.safe_load(open(self.vmdb_filename))
        self.boot = args.boot
        self.maybe_boot = args.maybe_boot
        self.verbose = args.verbose

        if args.arch:
            arch = args.arch
        else:
            arch = self.default_arch()
        self.qemu = QEMUS.get(arch)
        if self.boot and not self.qemu:
            sys.exit(
                f"Can't boot architecture {arch}: don't know of an emulator for it"
            )

        mklabel = self._step("mklabel")
        debootstrap = self._step("debootstrap")
        grub = self._step("grub")

        if args.debian_release is not None:
            debootstrap["debootstrap"] = args.debian_release
        if args.arch is not None:
            debootstrap["arch"] = args.arch
            debootstrap["include"] = debootstrap.get("include", []) + [
                KERNELS[args.arch]
            ]
        if args.mklabel is not None:
            mklabel["mklabel"] = args.mklabel
        if args.grub is not None:
            grub["grub"] = args.grub

    def _step(self, wanted):
        for item in self.vmdb["steps"]:
            if wanted in item:
                return item
        assert 0

    def debootstrap_release(self):
        for item in self.vmdb["steps"]:
            if "debootstrap" in item:
                return item["debootstrap"]
        assert 0

    def arch(self):
        for item in self.vmdb["steps"]:
            if "debootstrap" in item:
                return item.get("arch", self.default_arch())
        assert 0

    def default_arch(self):
        p = subprocess.run(
            ["dpkg", "--print-architecture"], check=True, capture_output=True
        )
        return p.stdout.decode().strip()

    def log(self):
        log = f"{self.debootstrap_release()}_{self.arch()}.log"
        return os.path.join(self.tarball_directory, log)

    def image(self):
        image = f"{self.debootstrap_release()}_{self.arch()}.img"
        return os.path.join(self.tarball_directory, image)

    def tarball(self):
        tarball = f"{self.debootstrap_release()}_{self.arch()}.tar.gz"
        return os.path.join(self.tarball_directory, tarball)

    def write_vmdb(self, filename):
        with open(filename, "w") as f:
            yaml.safe_dump(self.vmdb, stream=f, indent=4)


def run_vmdb2(config):
    print(f"building image {config.image()}")
    fd, vmdb = tempfile.mkstemp(dir=config.tarball_directory, suffix=".vmdb")
    config.write_vmdb(vmdb)
    os.close(fd)

    argv = [
        "./vmdb2",
        "--rootfs-tarball",
        config.tarball(),
        "--verbose",
        "--log",
        config.log(),
        "--output",
        config.image(),
        vmdb,
    ]
    subprocess.run(argv, check=True, capture_output=not config.verbose)

    os.remove(vmdb)
    print(f"built image {config.image()} OK")


def smoke_test(config):
    print(f"booting image {config.image()}")

    assert config.qemu
    tmp = tempfile.mkdtemp()
    qemu_sh = os.path.join(tmp, "run.sh")
    expect_txt = os.path.join(tmp, "expect.txt")

    qemu_script = f"""\
#!/bin/bash

set -euo pipefail
cd {tmp}
cp /usr/share/OVMF/OVMF_VARS.fd .
qemu-system-x86_64 \
  -m 1024 \
  -drive if=pflash,format=raw,unit=0,file=/usr/share/ovmf/OVMF.fd,readonly=on \
  -drive if=pflash,format=raw,unit=1,file=OVMF_VARS.fd \
  -drive format=raw,file="{config.image()}" \
  -nographic
"""

    expect_script = f"""\
set timeout 300
proc abort {{}} {{
    puts "ERROR ERROR\n"
    exit 1
}}
spawn {qemu_sh}
expect "login: "
send "root\n"
expect "# "
send "poweroff\r"
set timeout 5
expect {{
    "reboot: Power down" {{puts poweroffing\n}}
    eof abort
    timeout abort
}}
expect eof
"""

    with open(qemu_sh, "w") as f:
        f.write(qemu_script)
    os.chmod(qemu_sh, 0o755)

    with open(expect_txt, "w") as f:
        f.write(expect_script)

    p = subprocess.run(
        ["expect", "-d", expect_txt], check=False, capture_output=not config.verbose
    )
    shutil.rmtree(tmp)
    if p.returncode != 0:
        sys.stderr.write(p.stderr.decode())
        sys.exit(f"{config.image()} failed to boot")
    print(f"verified that {config.image()} boots OK")


def main():
    config = Config()
    if config.dump:
        config.write_vmdb("/dev/stdout")
    else:
        run_vmdb2(config)
        if config.boot:
            smoke_test(config)
        elif config.maybe_boot and config.qemu:
            smoke_test(config)


main()