#!/usr/bin/python3

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


MAX_CACHE_SIZE = 1024**3
OVMF_FD = "/usr/share/ovmf/OVMF.fd"
OVMF_VARS = "/usr/share/OVMF/OVMF_VARS.fd"


class BuildSpec:
    def __init__(self, filename):
        with open(filename) as f:
            o = yaml.safe_load(f)
        self.image = os.path.expanduser(o["image"])
        self.src = os.path.expanduser(o.get("src", "."))

        if "cache" in o:
            self.cache = os.path.expanduser(o["cache"])
        else:
            self.cache = None

        if "dependencies" in o:
            self.dependencies = os.path.expanduser(o["dependencies"])
        else:
            self.dependencies = None

        if "artifact" in o:
            self.artifact = os.path.expanduser(o["artifact"])
        else:
            self.artifact = None


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--verbose", action="store_true")
    p.add_argument("--log", required=True)
    p.add_argument("build_spec")
    return p.parse_args()


def extract_tar_from_tar(filename, output):
    N = 1024
    BLOCK = 10 * 1024

    with open(filename, "rb") as f:
        tar = tarfile.open(fileobj=f, mode="r:")
        while True:
            info = tar.next()
            if info is None:
                break

        # We've read to the end of the final tar file entry.
        num_bytes = tar.fileobj.tell()

        with open(output, "wb") as o:
            # Copy first num_bytes of input file to output file.

            f.seek(0)
            n = 0
            while n < num_bytes:
                data = f.read(N)
                if data is None:
                    break
                o.write(data)
                n += len(data)

            # Pad output file with zeroes until next block border.

            while (n % BLOCK) != 0:
                o.write(b"\0")
                n += 1


def main():
    args = parse_args()

    if args.verbose:
        level = logging.DEBUG
    else:
        level = logging.WARNING
    logging.basicConfig(
        level=level, stream=sys.stdout, format="%(levelname)s %(message)s"
    )
    logging.debug(f"args: {args}")

    try:
        logging.debug(f"loading build spec from {args.build_spec}")
        build = BuildSpec(args.build_spec)
        logging.debug(f"build: {build}")

        tmp = tempfile.mkdtemp()
        logging.debug(f"created {tmp}")

        image = os.path.join(tmp, "vm.qcow2")
        output_drive = os.path.join(tmp, "artifacts.tar")
        src_tar = os.path.join(tmp, "src.tar")
        vars = os.path.join(tmp, "vars.fd")

        logging.info(f"copying {build.image} to {image}")
        shutil.copyfile(build.image, image)

        logging.info(f"copying {OVMF_VARS}")
        shutil.copyfile(OVMF_FD, vars)

        logging.info(f"create tar of source tree {build.src}")
        subprocess.run(["tar", "-cf", src_tar, "-C", build.src, "."], check=True)

        logging.info(f"create output drive {output_drive}")
        subprocess.run(
            ["qemu-img", "create", "-q", "-f", "raw", output_drive, "1G"], check=True
        )

        cache_drive = os.path.join(tmp, "cache.tar")
        logging.info(f"using {cache_drive} as cache drive")
        tar = tarfile.open(name=cache_drive, mode="w:")
        if build.cache is None:
            tar.close()
        else:
            # Tar up the cache directory.
            tar.add(build.cache, arcname=".")
            tar.close()

            # Make the tar the max size.
            length = os.path.getsize(cache_drive)
            if length < MAX_CACHE_SIZE:
                os.truncate(cache_drive, MAX_CACHE_SIZE)
            else:
                logging.info(
                    "cache drive {cache_drive} is larger than max allowed, oh well"
                )

        deps_drive = os.path.join(tmp, "deps.tar")
        logging.info(f"using {deps_drive} as the dependencies drive")
        tar = tarfile.open(name=deps_drive, mode="w:")
        if build.dependencies is None:
            tar.close()
        else:
            # Tar up the dependencies directory.
            tar.add(build.dependencies, arcname=".")
            tar.close()

        logging.info("run build in VM")
        argv = [
            "kvm",
            "-m",
            "16384",
            "-smp",
            "cpus=4",
            "-drive",
            f"if=pflash,format=raw,unit=0,file={OVMF_FD},readonly=on",
            "-drive",
            f"if=pflash,format=raw,unit=1,file={vars}",
            "-drive",
            f"format=qcow2,if=virtio,file={image}",
            "-drive",
            f"format=raw,if=virtio,file={src_tar},readonly=on",
            "-drive",
            f"format=raw,if=virtio,file={output_drive}",
            "-drive",
            f"format=raw,if=virtio,file={cache_drive}",
            "-drive",
            f"format=raw,if=virtio,file={deps_drive},readonly=on",
            "-nodefaults",
            "-display",
            "none",
            "-chardev",
            "stdio,id=serial0",
            "-serial",
            "chardev:serial0",
        ]
        logging.debug(f"run: {argv}")
        with open(args.log, "wb") as f:
            p = subprocess.run(argv, stdout=f)
            if p.returncode != 0:
                sys.stderr.write(p.stderr.decode())
                raise Exception("qemu-system failed")

        if build.artifact is not None:
            logging.info("extract artifacts, if any")
            extract_tar_from_tar(output_drive, build.artifact)

        if build.cache is not None:
            logging.info(f"extract potentially updated cache from {cache_drive}")
            subprocess.run(
                ["tar", "-xf", cache_drive, "-C", build.cache],
                check=True,
            )

        logging.info("Build finished")
    except Exception as e:
        sys.stderr.write(f"{e}\n")
        sys.exit(1)
    finally:
        logging.debug(f"removing {tmp}")
        shutil.rmtree(tmp)


main()