#!/usr/bin/env python3

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


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, 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",
            "--network=network=default",
            "--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]))

    base = config["base_image"]
    vm = config["name"]
    img = config["image_file"]
    size = config["image_size"]
    pubkey = config["public_key"]
    memory = config.get("memory", 1024)
    cpus = config.get("cpus", 1)

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

    cloud_init_iso(iso, vm, pubkey)

    if os.path.exists(img):
        os.remove(img)
    shutil.copy(base, img)

    subprocess.check_call(["qemu-img", "resize", "-q", img, size])

    create_vm(vm, img, iso, memory=memory, cpus=cpus)
    wait_for_ssh(vm)
    os.remove(iso)


#    provision(vm, pubkey)


main()