#!/usr/bin/env python3

import os
import shutil
import socket
import subprocess
import sys
import tempfile
import time
import yaml


def msg(s):
    print(f"{s}")


def cloud_init_iso(iso, hostname, pubkey):
    tmp = tempfile.mkdtemp()

    with open(os.path.join(tmp, "meta-data"), "w") as f:
        f.write(
            f"""\
# Amazon EC2 style metadata
local-hostname: {hostname}
"""
        )

    with open(os.path.join(tmp, "user-data"), "w") as f:
        f.write(
            f"""\
#cloud-config
ssh_authorized_keys:
- {pubkey}
"""
        )

    subprocess.check_call(
        [
            "genisoimage",
            "-quiet",
            "-volid",
            "cidata",
            "-joliet",
            "-rock",
            "-output",
            iso,
            tmp,
        ]
    )
    shutil.rmtree(tmp)


def create_vm(vm, image, iso, network, memory=1024, cpus=1):
    subprocess.check_call(
        [
            "virt-install",
            "--name",
            vm,
            "--memory",
            str(memory),
            "--vcpus",
            str(cpus),
            f"--disk=path={image},cache=none",
            f"--disk=path={iso},readonly=on",
            f"--network={network}",
            "--connect",
            "qemu:///system",
            "--cpu=host-passthrough",
            "--os-variant=debian9",
            "--import",
            "--graphics=spice",
            "--noautoconsole",
            "--quiet",
        ]
    )


def wait_for_ssh(hostname):
    ssh_port = 22
    while True:
        time.sleep(5)
        try:
            conn = socket.create_connection((hostname, ssh_port), timeout=1)
        except Exception:
            continue
        conn.close()
        break


def provision(vm, pubkey):
    ssh_opts = [
        "ControlMaster=auto",
        "ControlPersist=60s",
        "StrictHostKeyChecking=accept-new",
        "UserKnownHostsFile=/dev/null",
    ]

    env = dict(os.environ)
    env["ANSIBLE_SSH_ARGS"] = " ".join(f"-o{opt}" for opt in ssh_opts)

    vars_file = {"user_pub": pubkey}

    fd, filename = tempfile.mkstemp()
    os.write(fd, yaml.safe_dump(vars_file).encode("UTF-8"))
    os.close(fd)

    argv = [
        "ansible-playbook",
        "-i",
        "hosts",
        "manager.yml",
        f"-eansible_ssh_host={vm}",
        f"-e@{filename}",
    ]
    subprocess.check_output(argv, env=env)

    os.remove(filename)


def main():
    config = yaml.safe_load(open(sys.argv[1]))
    vm = sys.argv[2]

    msg(f"creating VM {vm}")

    config = config[vm]
    base = os.path.expanduser(config["base_image"])
    img = os.path.expanduser(config["image_file"])
    size = config["image_size"]
    pubkey = os.path.expanduser(config["public_key"])
    memory = config.get("memory", 1024)
    cpus = config.get("cpus", 1)
    network = config.get("network", "network=default")

    memory = int(memory)
    cpus = int(cpus)
    pubkey = open(pubkey).read().rstrip()
    iso = f"{vm}.iso"

    msg(f"creating cloud-init config {iso}")
    cloud_init_iso(iso, vm, pubkey)

    if os.path.exists(img):
        msg(f"removing existing image {img}")
        os.remove(img)
    msg(f"creating new image {img} as a copy of {base}")
    shutil.copy(base, img)

    msg(f"resizing to {size}")
    subprocess.check_call(["qemu-img", "resize", "-q", img, size])

    msg(f"creating VM")
    create_vm(vm, img, iso, network, memory=memory, cpus=cpus)

    msg(f"waiting for SSH to be available")
    wait_for_ssh(vm)

    msg(f"removing cloud-init config")
    subprocess.check_output(
        ["virsh", "detach-disk", "--persistent", "--live", vm, "vdb"]
    )
    os.remove(iso)

    msg("OK")


main()