#!/usr/bin/python3

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


PROG = "ambient-build-debian-image"
DESC = "build virtual machine image with Debian for Ambient"
VERSION = "0.0"
CACHE = os.path.expanduser("~/.cache/ambient")
BASE_VMDB = "ambient.vmdb"
BASE_PLAYBOOK = "ambient-playbook-base.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(
        "--cache",
        action="store",
        help="set cache directory (default is %(default)s)",
        default=CACHE,
    )

    p.add_argument(
        "--playbook",
        action="store",
        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,
    )
    return p.parse_args()


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")
    extra_playbook_filename = os.path.join(tmp, "extra.yml")

    vmdb = load_yaml(BASE_VMDB)

    if args.playbook is not None:
        vmdb["steps"].append(
            {
                # note: the value MUST be the root file system tag in the VMDB file
                "ansible": "/",
                "playbook": args.playbook,
            }
        )
        shutil.copyfile(args.playbook, extra_playbook_filename)

    write_yaml(vmdb_filename, vmdb)
    shutil.copyfile(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.
    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.
    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}\nERROR: failed to convert image to QCOW2\n")
        sys.exit(1)


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

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

    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)


main()