#!/usr/bin/env node
// Vectrum runner: picks up workflow jobs from a Vectrum server and runs them on this machine.
//
//   node vectrum-runner.mjs --url https://vectrum.example --token vrn_…
//
// Options (or environment variables):
//   --url <url>          VECTRUM_URL           the Vectrum server
//   --token <token>      VECTRUM_RUNNER_TOKEN  the token shown when the runner was registered
//   --work-dir <dir>     VECTRUM_WORK_DIR      where jobs check out code (default ./_work)
//   --once                                     run one job, then exit
//   --keep-workspace                           leave each job's checkout behind
//
// Needs Node.js 18+ and git. Steps run as the user running this script, with its permissions:
// only register runners for repositories whose workflows you trust.

import { spawn, spawnSync } from "node:child_process";
import { mkdir, rm, writeFile } from "node:fs/promises";
import os from "node:os";
import path from "node:path";
import { StringDecoder } from "node:string_decoder";

const VERSION = "1.0.0";

function parseArgs(argv) {
  const opts = {};
  for (let i = 0; i < argv.length; i++) {
    const arg = argv[i];
    if (arg === "--once" || arg === "--keep-workspace") opts[arg.slice(2)] = true;
    else if (arg === "--help" || arg === "-h") opts.help = true;
    else if (arg.startsWith("--")) {
      const [key, inline] = arg.slice(2).split("=", 2);
      opts[key] = inline ?? argv[++i];
    } else fail(`Unknown argument: ${arg}`);
  }
  return opts;
}

function fail(message) {
  console.error(`vectrum-runner: ${message}`);
  process.exit(1);
}

const opts = parseArgs(process.argv.slice(2));
if (opts.help) {
  console.log("Usage: node vectrum-runner.mjs --url <server> --token <vrn_…> [--work-dir <dir>] [--once] [--keep-workspace]");
  process.exit(0);
}
const SERVER = (opts.url ?? process.env.VECTRUM_URL ?? "").replace(/\/+$/, "");
const TOKEN = opts.token ?? process.env.VECTRUM_RUNNER_TOKEN ?? "";
const WORK_DIR = path.resolve(opts["work-dir"] ?? process.env.VECTRUM_WORK_DIR ?? "_work");
if (!SERVER) fail("pass the server with --url (or VECTRUM_URL)");
if (!TOKEN.startsWith("vrn_")) fail("pass the runner token with --token (or VECTRUM_RUNNER_TOKEN)");

const OS = { linux: "linux", darwin: "macos", win32: "windows" }[process.platform] ?? process.platform;
const ARCH = { x64: "x64", arm64: "arm64", arm: "arm", ia32: "x86" }[process.arch] ?? process.arch;
const IS_WINDOWS = process.platform === "win32";

const sleep = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
const stamp = () => new Date().toISOString().replace("T", " ").slice(0, 19);
const say = (message) => console.log(`${stamp()}  ${message}`);

class ApiError extends Error {
  constructor(message, status) {
    super(message);
    this.status = status;
  }
}

async function api(pathname, body, signal) {
  const res = await fetch(`${SERVER}${pathname}`, {
    method: "POST",
    headers: { authorization: `Bearer ${TOKEN}`, "content-type": "application/json" },
    body: JSON.stringify(body),
    signal,
  });
  if (res.status === 204) return null;
  const data = await res.json().catch(() => ({}));
  if (!res.ok) throw new ApiError(data.error ?? `HTTP ${res.status}`, res.status);
  return data;
}

// ---------------------------------------------------------------------------
// Reporting progress

/** Batches step changes and output, sending them about once a second (and a heartbeat when quiet). */
class Reporter {
  constructor(jobId) {
    this.jobId = jobId;
    this.steps = new Map();
    this.logs = [];
    this.cancelled = false;
    this.lost = false;
    this.lastSent = 0;
    this.sending = Promise.resolve();
    this.timer = setInterval(() => {
      if (this.logs.length || this.steps.size || Date.now() - this.lastSent > 15_000) void this.flush();
    }, 1000);
  }

  step(index, status, exitCode) {
    this.steps.set(index, { index, status, exitCode });
  }

  log(step, text) {
    const last = this.logs[this.logs.length - 1];
    if (last && last.step === step && last.text.length < 64 * 1024) last.text += text;
    else this.logs.push({ step, text });
  }

  /** Sends what is pending; calls are queued so updates arrive in order. */
  flush(extra = {}) {
    this.sending = this.sending.then(async () => {
      if (this.lost) return;
      const body = { steps: [...this.steps.values()], logs: this.logs, ...extra };
      this.steps = new Map();
      this.logs = [];
      for (let attempt = 0; ; attempt++) {
        try {
          const res = await api(`/api/runner/jobs/${this.jobId}`, body);
          this.lastSent = Date.now();
          if (res?.cancelled) this.cancelled = true;
          return;
        } catch (e) {
          if (e instanceof ApiError && (e.status === 409 || e.status === 401)) {
            // The server gave the job up (or the runner was removed): stop working on it.
            say(`  ${e.message}`);
            this.lost = true;
            this.cancelled = true;
            return;
          }
          if (attempt >= 20) {
            this.lost = true;
            this.cancelled = true;
            return;
          }
          await sleep(Math.min(1000 * 2 ** attempt, 15_000));
        }
      }
    });
    return this.sending;
  }

  async finish(status) {
    clearInterval(this.timer);
    await this.flush({ status });
  }
}

// ---------------------------------------------------------------------------
// Running commands

/** Stops a process and everything it started. */
function killTree(child) {
  if (child.exitCode !== null || child.signalCode !== null) return;
  if (IS_WINDOWS) {
    spawnSync("taskkill", ["/pid", String(child.pid), "/T", "/F"], { stdio: "ignore" });
    return;
  }
  try {
    process.kill(-child.pid, "SIGTERM");
  } catch {
    child.kill("SIGTERM");
  }
  setTimeout(() => {
    try {
      process.kill(-child.pid, "SIGKILL");
    } catch {
      // Already gone.
    }
  }, 5000).unref();
}

/**
 * Runs a command, streaming its output to `onOutput`. Resolves with its exit code, or with
 * "cancelled" / "timeout" if `shouldStop` said so or `timeoutMs` passed first.
 */
function execute(command, args, { cwd, env, onOutput, shouldStop, timeoutMs }) {
  return new Promise((resolve) => {
    let child;
    try {
      child = spawn(command, args, { cwd, env, detached: !IS_WINDOWS, stdio: ["ignore", "pipe", "pipe"], windowsHide: true });
    } catch (e) {
      onOutput(`${e.message}\n`);
      resolve(127);
      return;
    }
    let stopped = null;
    const decoders = [new StringDecoder("utf8"), new StringDecoder("utf8")];
    child.stdout.on("data", (d) => onOutput(decoders[0].write(d)));
    child.stderr.on("data", (d) => onOutput(decoders[1].write(d)));
    const watch = setInterval(() => {
      if (!stopped && shouldStop()) {
        stopped = "cancelled";
        killTree(child);
      }
    }, 500);
    const timer = timeoutMs
      ? setTimeout(() => {
          stopped = "timeout";
          killTree(child);
        }, timeoutMs)
      : null;
    child.on("error", (e) => onOutput(`${e.message}\n`));
    child.on("close", (code, signal) => {
      clearInterval(watch);
      if (timer) clearTimeout(timer);
      onOutput(decoders[0].end() + decoders[1].end());
      resolve(stopped ?? (code ?? (signal ? 128 : 1)));
    });
  });
}

const hasCommand = (cmd) => !spawnSync(cmd, IS_WINDOWS ? ["-NoProfile", "-Command", "exit 0"] : ["-c", "true"], { stdio: "ignore" }).error;
const DEFAULT_SHELL = IS_WINDOWS ? (hasCommand("pwsh") ? "pwsh" : "powershell") : hasCommand("bash") ? "bash" : "sh";

/** How a step's script file is run, by its `shell:` (GitHub's conventions). */
function shellCommand(shell, file) {
  switch (shell) {
    case "bash":
      return { ext: ".sh", command: "bash", args: ["--noprofile", "--norc", "-eo", "pipefail", file] };
    case "sh":
      return { ext: ".sh", command: "sh", args: ["-e", file] };
    case "pwsh":
    case "powershell":
      return { ext: ".ps1", command: shell, args: ["-NoProfile", "-NonInteractive", "-Command", `$ErrorActionPreference = 'stop'; & '${file}'; if ((Test-Path -LiteralPath variable:\\LASTEXITCODE)) { exit $LASTEXITCODE }`] };
    case "cmd":
      return { ext: ".cmd", command: "cmd", args: ["/D", "/E:ON", "/V:OFF", "/S", "/C", `CALL "${file}"`] };
    case "python":
      return { ext: ".py", command: IS_WINDOWS ? "python" : "python3", args: [file] };
    case "node":
      return { ext: ".mjs", command: "node", args: [file] };
    default: {
      // A custom command line, with {0} where the script file goes: "perl {0}".
      if (!shell.includes("{0}")) return null;
      const [command, ...args] = shell.trim().split(/\s+/).map((part) => part.replace("{0}", file));
      return { ext: "", command, args };
    }
  }
}

// ---------------------------------------------------------------------------
// Jobs

async function checkout(job, workspace, reporter, shouldStop) {
  const auth = Buffer.from(`x-access-token:${job.token}`).toString("base64");
  const out = (text) => reporter.log(0, text);
  const git = (args) => {
    out(`$ git ${args.map((a) => (a.startsWith("http.extraHeader=") ? "http.extraHeader=***" : a)).join(" ")}\n`);
    return execute("git", args, { cwd: workspace, env: { ...process.env, GIT_TERMINAL_PROMPT: "0" }, onOutput: out, shouldStop });
  };
  out(`Checking out ${job.repository} at ${job.sha.slice(0, 7)} (${job.ref})\n`);
  const init = await git(["init", "--quiet"]);
  if (init !== 0) return init;
  // The token goes in a header for this one command, so it isn't stored in .git/config.
  const fetched = await git([
    "-c", `http.extraHeader=Authorization: Basic ${auth}`,
    "-c", "protocol.version=2",
    "fetch", "--no-tags", "--depth=1", "--quiet", job.cloneUrl, job.sha,
  ]);
  if (fetched !== 0) return fetched;
  return git(["checkout", "--quiet", "--detach", "FETCH_HEAD"]);
}

async function runJob(job) {
  const workspace = path.join(WORK_DIR, String(job.id));
  const scripts = path.join(WORK_DIR, `${job.id}-scripts`);
  const reporter = new Reporter(job.id);
  say(`Running ${job.repository} ${job.env.VECTRUM_WORKFLOW} #${job.env.VECTRUM_RUN_NUMBER} · ${job.env.VECTRUM_JOB}`);
  let status;
  try {
    status = await runSteps(job, workspace, scripts, reporter);
  } catch (e) {
    reporter.log(0, `\nThe runner failed: ${e.stack ?? e.message}\n`);
    status = "failure";
  }
  await reporter.finish(status);
  say(`  ${status === "success" ? "Succeeded" : status === "failure" ? "Failed" : "Cancelled"}`);
  await rm(scripts, { recursive: true, force: true }).catch(() => {});
  if (!opts["keep-workspace"]) await rm(workspace, { recursive: true, force: true }).catch(() => {});
}

/** Checks out the code and runs the steps; returns the job's result. */
async function runSteps(job, workspace, scripts, reporter) {
  const shouldStop = () => reporter.cancelled || stopping;
  const deadline = Date.now() + job.timeoutMinutes * 60_000;
  let failed = false;
  let timedOut = false;

  await rm(workspace, { recursive: true, force: true });
  await mkdir(workspace, { recursive: true });
  await mkdir(scripts, { recursive: true });

  reporter.step(0, "in_progress");
  const checkedOut = await checkout(job, workspace, reporter, shouldStop);
  if (checkedOut !== 0) {
    failed = true;
    reporter.log(0, checkedOut === "cancelled" ? "Cancelled.\n" : `Checkout failed (${checkedOut}).\n`);
  }
  reporter.step(0, checkedOut === 0 ? "success" : checkedOut === "cancelled" ? "cancelled" : "failure", typeof checkedOut === "number" ? checkedOut : null);

  for (const step of job.steps) {
    const cancelled = shouldStop() || timedOut;
    const run = step.if === "always" || (cancelled ? false : step.if === "failure" ? failed : !failed);
    if (!run) {
      reporter.step(step.index, "skipped");
      continue;
    }
    reporter.step(step.index, "in_progress");
    const out = (text) => reporter.log(step.index, text);

    const shell = shellCommand(step.shell ?? DEFAULT_SHELL, "");
    if (!shell) {
      out(`Unknown shell “${step.shell}”. Use bash, sh, pwsh, powershell, cmd, python, node, or a command with {0}.\n`);
      reporter.step(step.index, "failure");
      if (!step.continueOnError) failed = true;
      continue;
    }
    const file = path.join(scripts, `step-${step.index}${shell.ext}`);
    await writeFile(file, step.run);
    const { command, args } = shellCommand(step.shell ?? DEFAULT_SHELL, file);
    const cwd = step.workingDirectory ? path.resolve(workspace, step.workingDirectory) : workspace;
    const stepTimeout = step.timeoutMinutes ? step.timeoutMinutes * 60_000 : Infinity;
    const remaining = deadline - Date.now();

    const result = await execute(command, args, {
      cwd,
      env: { ...process.env, ...job.env, ...step.env, VECTRUM_WORKSPACE: workspace, RUNNER_OS: OS, RUNNER_ARCH: ARCH },
      onOutput: out,
      // always() steps run on after a cancel; only stopping the runner interrupts them.
      shouldStop: step.if === "always" ? () => stopping : shouldStop,
      timeoutMs: Math.max(1000, Math.min(stepTimeout, remaining)),
    });

    if (result === 0) {
      reporter.step(step.index, "success", 0);
    } else if (result === "cancelled") {
      out("\nCancelled.\n");
      reporter.step(step.index, "cancelled");
    } else {
      if (result === "timeout") {
        out(`\nThe step ran out of time.\n`);
        if (Date.now() >= deadline - 1000) timedOut = true;
      } else {
        out(`\nProcess exited with code ${result}.\n`);
      }
      reporter.step(step.index, "failure", typeof result === "number" ? result : null);
      if (!step.continueOnError) failed = true;
    }
  }

  if (timedOut) reporter.log(job.steps.at(-1)?.index ?? 0, `\nThe job ran longer than its ${job.timeoutMinutes} minute limit.\n`);
  return reporter.cancelled || stopping ? "cancelled" : failed || timedOut ? "failure" : "success";
}

// ---------------------------------------------------------------------------
// Main loop

let stopping = false;
let busy = false;
process.on("SIGINT", onSignal);
process.on("SIGTERM", onSignal);

function onSignal() {
  if (stopping || !busy) process.exit(130);
  stopping = true;
  say("Stopping: cancelling the current job (press Ctrl-C again to quit now)…");
}

const info = { hostname: os.hostname(), os: OS, arch: ARCH, version: VERSION };
say(`Vectrum runner ${VERSION} on ${info.hostname} (${OS}/${ARCH}), waiting for jobs from ${SERVER}`);
let failures = 0;

while (!stopping) {
  let job;
  try {
    job = (await api("/api/runner/claim", info, AbortSignal.timeout(60_000)))?.job;
    failures = 0;
  } catch (e) {
    if (e instanceof ApiError && e.status === 401) fail(e.message);
    failures++;
    say(`Can't reach ${SERVER}: ${e.cause?.message ?? e.message}. Retrying…`);
    await sleep(Math.min(2000 * failures, 30_000));
    continue;
  }
  if (!job) continue;
  busy = true;
  await runJob(job);
  busy = false;
  if (opts.once) break;
}