#!/usr/bin/python3

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


PROG = "ambient-build-vm"
DESC = "build virtual machine image with Debian for Ambient"
VERSION = "0.1.0"
CACHE = os.path.expanduser("~/.cache/ambient")
BASE_VMDB = "base.vmdb"
BASE_PLAYBOOK = "playbook.yml"


def parse_args():
    p = argparse.ArgumentParser(prog=PROG, description=DESC)

    p.add_argument(
        "--version", action="version", version=VERSION, help="show version of program"
    )

    p.add_argument(
        "--data",
        default=f"/usr/share/{PROG}",
        help="look for data files in DATA directory",
    )

    p.add_argument(
        "--cache",
        action="store",
        help="set cache directory (default is %(default)s)",
        default=CACHE,
    )

    p.add_argument(
        "--playbook",
        action="append",
        help="extra Ansible playbook to use with vmdb to build image (default is none)",
    )

    p.add_argument(
        "--image",
        action="store",
        help="filename for resulting image (required)",
        required=True,
    )

    p.add_argument(
        "--debian",
        action="store",
        help="install Debian version (codename) instead of the default (%default)",
        default="bookworm",
    )

    return p.parse_args()


def join(args, relative):
    return os.path.join(args.data, relative)


def load_yaml(filename):
    with open(filename) as f:
        return yaml.safe_load(f)


def write_yaml(filename, obj):
    with open(filename, "w") as f:
        yaml.safe_dump(obj, stream=f, indent=4)


def prepare_build_files(tmp, args):
    vmdb_filename = os.path.join(tmp, "image.vmdb")
    playbook_filename = os.path.join(tmp, "playbook.yml")

    vmdb = load_yaml(join(args, BASE_VMDB))

    for step in vmdb["steps"]:
        if "debootstrap" in step:
            step["debootstrap"] = args.debian

    for playbook in args.playbook or []:
        vmdb["steps"].append(
            {
                # note: the value MUST be the root file system tag in the VMDB file
                "ansible": "/",
                "playbook": playbook,
            }
        )
        shutil.copyfile(playbook, os.path.join(tmp, playbook))

    write_yaml(vmdb_filename, vmdb)
    shutil.copyfile(join(args, BASE_PLAYBOOK), playbook_filename)

    return vmdb_filename


def build_image(tmp, vmdb, tarball, image):
    raw = os.path.join(tmp, "image.img")
    log = os.path.join(tmp, "log")

    # Build image.
    logging.info(f"Building RAW image {raw} with vmdb2, using {tarball}")
    p = subprocess.run(
        ["vmdb2", vmdb, "--output", raw, "--rootfs-tarball", tarball, "--log", log],
        capture_output=True,
    )
    if p.returncode != 0:
        with open(log, "r") as f:
            log = f.read()
        sys.stderr.write(f"{log}\nERROR: failed to build image\n")
        sys.exit(1)

    # Convert built image from raw to QCOW2 format.
    logging.info(f"Converting RAW image to {image}")
    p = subprocess.run(
        ["qemu-img", "convert", "-f", "raw", "-O", "qcow2", raw, image],
        capture_output=True,
    )
    if p.returncode != 0:
        sys.stderr.write(
            f"{p.stderr.decode()}\nERROR: failed to convert image to QCOW2\n"
        )
        sys.exit(1)


def main():
    logging.basicConfig(
        level=logging.DEBUG, stream=sys.stdout, format="%(levelname)s %(message)s"
    )
    logging.info(f"{PROG} {VERSION} starts")

    args = parse_args()
    tarball = os.path.join(args.cache, "ambient.tar.gz")
    tmp = tempfile.mkdtemp()

    logging.info("Creating build files")
    try:
        vmdb = prepare_build_files(tmp, args)
    except Exception as e:
        sys.stderr.write(f"ERROR: failed to create build files: {e}\n")
        shutil.rmtree(tmp)
        sys.exit(1)

    logging.info(f"Build image {args.image}")
    try:
        build_image(tmp, vmdb, tarball, args.image)
    except Exception as e:
        sys.stderr.write(f"ERROR: failed to build image: {e}\n")
        shutil.rmtree(tmp)
        sys.exit(1)

    logging.info("All good")


main()