#!/usr/bin/env python3
"""Provision a large virtual machine on Digital Ocean and run a nix remote
build, then destroy the VM.

I tested this like so:

   $ ./nix-build-omatic "'<nixpkgs>'" -A ocaml

The special quoting is necessary until I setup an argparse CLI, currently it
just passes all args through to the nix-build call. But it works!

"""

import json
import os
import sys
import subprocess
import time

def noti(msg):
    print(f"nix-build-omatic: {msg}")

class Builder:
    def __init__(self):
        # TODO: generate a temporary privkey, and use that instead
        self.keys = [k["id"] for k in json.loads(subprocess.check_output(["doctl", "compute", "ssh-key", "list", "--output", "json"]).decode('utf-8'))]
        self.privkey_path = os.path.expanduser("~/.ssh/id_rsa")
        self.max_jobs = 32
        self.ip = None

    def __enter__(self):
        "Startup a droplet."
        noti("preparing droplet...")
        res = self.droplet(
            "create",
            # this is my custom nixos image, I had to download from nixos.org and then
            # upload it to DO. There's gotta be a better way to do this..
            "--image", "121034270",
            "--size", "c2-32vcpu-64gb", # TODO: set this in argparse, match '32vcpu' with 'max_jobs'
            "--region", "nyc3",
            "--tag-name", "nix-build-omatic",
            "--ssh-keys", ",".join([str(k) for k in self.keys]),
            "nix-build-omatic"
        )
        self.id = str(res["id"])
        self.active = False
        while not self.active:
            res = self.droplet("get", self.id)
            if res["status"] == "active":
                self.active = True
                noti("droplet ready!")
            else:
                continue
        while self.ip is None:
            res = self.droplet("get", self.id)
            self.ip = [n["ip_address"] for n in res["networks"]["v4"] if n["type"] == "public" ][0]
        self.url = f"ssh://root@{self.ip}"
        self.do_stupid_handshake()
        return self

    def __exit__(self, type, value, traceback):
        "Destroy the droplet."
        self.droplet("delete", self.id, "--force", parse=False)
        noti("droplet destroyed")

    def droplet(self, *args, parse=True):
        "Run a `doctl compute droplet *args` command"
        proc = subprocess.run(
            ["doctl", "compute", "droplet", *args, "--output", "json"],
            capture_output=True
        )
        stderr = proc.stderr.decode('utf-8')
        stdout = proc.stdout.decode('utf-8')
        if proc.returncode > 0:
            noti("stderr:", stderr)
            noti("stdout:", stdout)
        if parse:
            return json.loads(stdout)[0]
        else:
            return stdout

    def do_stupid_handshake(self):
        """
        this whole thing is doing a key handshake, which shouldn't be necessary!
        i should just generate the keys on my end and set them via the API! or
        retrieve them via the API! i'm already authenticated! wtf digital ocean
        """
        noti("exchanging keys...")
        keyscan = None
        while keyscan is None:
            cmd = f"ssh-keyscan -t ed25519 {self.ip}"
            proc = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            stderr = proc.stderr.decode('utf-8')
            stdout = proc.stdout.decode('utf-8')
            if proc.returncode > 0:
                continue
            elif stdout == "":
                continue
            else:
                keyscan = stdout
        with open(os.path.expanduser("~/.ssh/known_hosts"), 'a') as known_hosts:
            known_hosts.write(keyscan)
        self.host_pubkey = subprocess.check_output(f"ssh root@{self.ip} cat /etc/ssh/ssh_host_ed25519_key.pub | base64 -w0", shell=True).decode('utf-8')


if __name__ == '__main__':
    # TODO: use argparse, properly quote '<nixpkgs>' and so on
    args = sys.argv[1:]
    start_time = time.time()

    with Builder() as b:
        cmd = " ".join([
            "nix-build",
            *args,
            # TODO: obv get rid of this when I'm not testing...
            "--option", "substitute", "false",
            # TODO: make this --builders line structured data somehow
            "--builders", f"'{b.url} x86_64-linux {b.privkey_path} {b.max_jobs} - - - {b.host_pubkey}'"
        ])
        noti(f"running: {cmd}")
        subprocess.check_call(cmd, shell=True)

    duration = time.time() - start_time
    duration_human = time.strftime('%H:%M:%S', duration)
    # the 32 vCPU machine costs $1/hr, so convert that to cost per second
    cost = (1/60/60)*duration
    print(f"duration: {duration_human}")
    print(f"estimated cost: {cost:0.2f} USD")