#!/usr/bin/env python3
"""Run each Buildkite job on a fresh exe.dev sandbox.

From https://exe.dev/deploy/buildkite. Python 3 standard library only.

A custom stack on Buildkite's Stacks API (https://buildkite.com/docs/apis/agent-api/stacks).
Every POLL_SECONDS, for each queue in QUEUES, it:

  * lists the jobs waiting in the queue and reserves as many as there is room
    for (MAX_AGENTS job VMs at once);
  * for each one, copies BASE_VM into a sandbox, asks Buildkite for a job
    acquisition token, and starts an agent on the copy that runs that job only;
  * deletes the copies whose jobs have finished, been cancelled, or never
    started.

The cluster agent token and the exe.dev API key stay in the integrations
AGENT_API and EXE_API; job VMs get a token for their one job. It only touches
VMs whose names start with JOB_PREFIX, and records each one's job in its
exe.dev comment, so it picks up where it left off after a restart.
"""
import base64, json, os, secrets, signal, socket, sys, threading, time
import urllib.error, urllib.request
from datetime import datetime

ENV = os.environ
AGENT_API = ENV["AGENT_API"].rstrip("/")  # the integration, e.g. https://bk-agent.int.exe.xyz
EXE_API = ENV["EXE_API"].rstrip("/")      # e.g. https://bk-exe.int.exe.xyz
BASE_VM = ENV["BASE_VM"]
# Queue keys and the vCPUs of their job VMs; each vCPU comes with 2 GiB.
QUEUES = dict(q.split("=") for q in ENV.get("QUEUES", "exe=2").split())
STACK_KEY = ENV.get("STACK_KEY", socket.gethostname())
JOB_PREFIX = ENV.get("JOB_PREFIX", socket.gethostname() + "-job-")
MAX_AGENTS = int(ENV.get("MAX_AGENTS", "5"))
START_SECONDS = int(ENV.get("START_SECONDS", "300"))  # to pick up its job
MAX_HOURS = float(ENV.get("MAX_HOURS", "6"))
POLL_SECONDS = float(ENV.get("POLL_SECONDS", "2"))
ACTIVE = {"scheduled", "reserved", "assigned", "accepted", "running",
          "canceling", "timing_out"}


def log(*a):
    print(time.strftime("%H:%M:%S"), *a, flush=True)


def http(method, url, body=None, accept="application/json"):
    data = None if body is None else (body if isinstance(body, bytes) else json.dumps(body).encode())
    req = urllib.request.Request(url, data=data, method=method, headers={
        "Content-Type": "application/json", "Accept": accept,
        "User-Agent": "exe-bk-dispatcher"})
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            raw = r.read().decode()
    except urllib.error.HTTPError as e:
        # Keep the reason exe.dev or Buildkite gives, for the log and the job.
        try:
            why = json.loads(e.read()).get("error") or e.reason
        except Exception:
            why = e.reason
        raise urllib.error.HTTPError(url, e.code, why, e.headers, None)
    lines = raw.strip().splitlines()
    return json.loads(lines[-1]) if lines else None  # exe.dev may add a line before the result


def bk(method, path, body=None):
    """Buildkite's Stacks API, through the integration that adds the token."""
    for attempt in range(5):
        try:
            return http(method, f"{AGENT_API}/v3/stacks{path}", body)
        except urllib.error.HTTPError as e:
            if e.code != 429 and e.code < 500:
                raise
            time.sleep(float(e.headers.get("Retry-After") or 1 + attempt))
    return http(method, f"{AGENT_API}/v3/stacks{path}", body)


def exe(cmd):
    """Run an exe.dev command over the HTTPS API; returns its JSON."""
    return http("POST", EXE_API + "/exec", cmd.encode())


def exe_ssh(vm, command):
    """Run a command on a VM through the API; returns its exit code."""
    req = urllib.request.Request(EXE_API + "/exec", method="POST",
                                 data=f"ssh {vm} '{command}'".encode(),
                                 headers={"Accept": "application/x-ndjson",
                                          "User-Agent": "exe-bk-dispatcher"})
    code = None
    with urllib.request.urlopen(req, timeout=60) as r:
        for line in r:
            code = json.loads(line).get("exit_code", code)
    return code


lock = threading.Lock()
vms = {}      # our job VMs: name -> [job UUID or None, created (epoch seconds)]
starting = set()  # job VMs being launched right now
listed = 0


def sync_vms():
    """Refresh vms from exe.dev: at startup, then every minute."""
    global listed
    found = {}
    for v in exe("ls --json")["vms"]:
        if v["vm_name"].startswith(JOB_PREFIX):
            t = datetime.fromisoformat(v["created_at"].replace("Z", "+00:00"))
            words = (v.get("comment") or "").split()
            job = words[1] if len(words) == 2 and words[0] == "buildkite" else None
            found[v["vm_name"]] = [job, t.timestamp()]
    with lock:
        for name in starting:
            found.setdefault(name, vms.get(name))
        vms.clear()
        vms.update(found)
    listed = time.time()


def remove(name, why):
    try:
        # Let the agent disconnect from Buildkite before its VM goes, or
        # Buildkite lists it as connected until it gives up on it.
        exe_ssh(name, "sudo systemctl stop buildkite-job")
    except Exception:
        pass
    try:
        exe(f"rm {name}")
        log(f"deleted {name} ({why})")
    except Exception as e:
        log(f"rm {name}: {e}")  # still listed by exe.dev? the next sync finds it again
    with lock:
        vms.pop(name, None)


def fail(job, detail):
    try:
        bk("POST", f"/{STACK_KEY}/jobs/{job}/finish", {"exit_status": -1, "detail": detail})
    except Exception as e:
        log(f"finish {job}: {e}")


def launch(queue, job):
    name = JOB_PREFIX + secrets.token_hex(4)
    cpus = int(QUEUES[queue])
    t0 = time.time()
    with lock:
        starting.add(name)
        vms[name] = [job, t0]
    try:
        exe(f"cp {BASE_VM} {name} --sandbox --cpu={cpus} --memory={2 * cpus}GB")
        exe(f"comment {name} buildkite {job}")
        for _ in range(30):  # the copy takes a moment to answer
            try:
                if exe_ssh(name, "true") == 0:
                    break
            except urllib.error.URLError:
                pass
            time.sleep(1)
        else:
            raise RuntimeError("the VM did not answer")
        # Only now, with somewhere to run it, ask for the job's token.
        res = bk("POST", f"/{STACK_KEY}/job-acquisition-tokens",
                 {"job_uuids": [job], "token_lifetime_seconds": 600})
        jat = next(t["job_acquisition_token"] for t in res["job_acquisition_tokens"]
                   if t["job_uuid"] == job)
        env = (f"BUILDKITE_AGENT_TOKEN={jat}\nBUILDKITE_AGENT_ACQUIRE_JOB={job}\n"
               f"BUILDKITE_AGENT_NAME={name}\nBUILDKITE_AGENT_TAGS=queue={queue}\n")
        b64 = base64.b64encode(env.encode()).decode()
        start = (f"echo {b64} | base64 -d | sudo install -m 600 /dev/stdin /run/buildkite-job.env"
                 " && sudo systemctl start --no-block buildkite-job")
        if exe_ssh(name, start) != 0:
            raise RuntimeError("could not start the agent")
        log(f"started {name} for {queue} job {job} in {time.time() - t0:.1f}s")
    except Exception as e:
        log(f"launch {name} for {job} failed: {e}")
        fail(job, f"exe.dev: could not start {name}: {e}"[:4000])
        remove(name, "launch failed")
    finally:
        with lock:
            starting.discard(name)


def reconcile():
    if time.time() - listed > 60:
        sync_vms()
    with lock:
        mine = {n: v for n, v in vms.items() if n not in starting}
        alive = len(vms) - len(mine)
    # Delete the VMs whose jobs are over, or that never got going.
    jobs = [j for j, _ in mine.values() if j]
    states = bk("POST", f"/{STACK_KEY}/jobs/get-states", {"job_uuids": jobs})["states"] if jobs else {}
    gone = []
    for name, (job, created) in mine.items():
        age = time.time() - created
        state = states.get(job, "unknown")
        if job is None:
            if age > 120:
                gone.append((name, "no job"))
            else:
                alive += 1
        elif state not in ACTIVE:
            gone.append((name, f"job {state}"))
        elif state in ("scheduled", "reserved") and age > START_SECONDS:
            fail(job, f"exe.dev: the agent on {name} did not pick up the job")
            gone.append((name, "agent did not start"))
        elif age > MAX_HOURS * 3600:
            gone.append((name, f"older than {MAX_HOURS:g} hours"))
        else:
            alive += 1
    threads = [threading.Thread(target=remove, args=g) for g in gone]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    # Start VMs for waiting jobs, oldest first, while there is room.
    for queue in QUEUES:
        room = MAX_AGENTS - alive
        if room <= 0:
            return
        res = bk("GET", f"/{STACK_KEY}/scheduled-jobs?queue_key={queue}&limit={room}")
        if res["cluster_queue"].get("dispatch_paused") or not res["jobs"]:
            continue
        want = [j["id"] for j in res["jobs"]][:room]
        got = bk("PUT", f"/{STACK_KEY}/scheduled-jobs/batch-reserve",
                 {"job_uuids": want, "reservation_expiry_seconds": START_SECONDS})["reserved"]
        for job in got:
            threading.Thread(target=launch, args=(queue, job), daemon=True).start()
            alive += 1


def register():
    # One stack serves every queue in QUEUES; Buildkite doesn't hold it to the first.
    bk("POST", "/register", {"key": STACK_KEY, "type": "custom", "queue_key": next(iter(QUEUES)),
                             "metadata": {"base_vm": BASE_VM, "queues": " ".join(QUEUES)}})


def stop(*_):
    try:
        bk("POST", f"/{STACK_KEY}/deregister")
    except Exception as e:
        log(f"deregister: {e!r}")
    sys.exit(0)


if __name__ == "__main__":
    signal.signal(signal.SIGTERM, stop)
    log(f"stack {STACK_KEY}: queues {' '.join(f'{q} ({c} vCPU)' for q, c in QUEUES.items())}, "
        f"copies of {BASE_VM}, up to {MAX_AGENTS} at once")
    while True:
        try:
            register()
            break
        except Exception as e:
            log(f"register: {e!r}")
            time.sleep(10)
    while True:
        try:
            reconcile()
        except Exception as e:
            log(f"reconcile: {e!r}")
        time.sleep(POLL_SECONDS)