//! Runner updates, the runner's side (docs/jetstream.md §8.2, #85): check

//! a new `flow` build, stage it beside the running one, swap it in keeping

//! the old one as `flow.prev`, and have the running runner leave with

//! notice so its unit starts the new build. The new build proves itself by

//! reaching the hub; three failed starts, or ten minutes without the hub,

//! put `flow.prev` back.

use anyhow::{bail, Context, Result};

use serde::{Deserialize, Serialize};

use std::path::{Path, PathBuf};

use std::time::Duration;

/// The staged build, beside the installed one.

pub const NEW: &str = "flow.new";

/// The build an update replaced, kept to go back to.

pub const PREV: &str = "flow.prev";

/// In the runner's home: the update not yet proven.

const PENDING: &str = "update.json";

/// In the runner's home: every update and rollback, a JSON line each.

const LOG: &str = "updates.log";

/// In the runner's home: the running runner's process id.

const PID: &str = "runner.pid";

/// Starts of an unproven build after which the previous one comes back.

pub const MAX_STARTS: u32 = 3;

/// How long a new build has to reach the hub.

pub const PROVE_WITHIN: Duration = Duration::from_secs(10 * 60);

/// An update swapped in and not yet proven.

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]

pub struct Pending {

    pub from: String,

    pub to: String,

    /// The installed binary the update replaced.

    pub exe: PathBuf,

    pub at_ms: u64,

    /// How long the running runner had to finish what it held.

    pub notice_secs: u64,

    /// Starts of the new build so far.

    pub starts: u32,

}

/// What [`install`] did, or would do.

#[derive(Debug, PartialEq, Eq)]

pub struct Installed {

    pub from: String,

    pub to: String,

    pub exe: PathBuf,

    /// The running runner told to leave, by process id.

    pub told: Option<u32>,

}

/// The version a `flow --version` line names: `flow 0.3.1 (build)` → `0.3.1`.

pub fn parse_version(line: &str) -> Option<String> {

    let v = line.trim().strip_prefix("flow ")?.split_whitespace().next()?;

    v.split('.').all(|p| !p.is_empty() && p.bytes().all(|b| b.is_ascii_digit())).then(|| v.to_string())

}

/// Whether version `a` is older than `b`, part by part.

pub fn older(a: &str, b: &str) -> bool {

    let parts = |v: &str| v.split('.').map(|p| p.parse::<u64>().unwrap_or(0)).collect::<Vec<_>>();

    parts(a) < parts(b)

}

/// The ELF machine number of this build's architecture.

fn machine_here() -> Option<u16> {

    Some(match std::env::consts::ARCH {

        "x86_64" => 62,

        "aarch64" => 183,

        "riscv64" => 243,

        _ => return None,

    })

}

/// Refuse a file that is not a 64-bit little-endian ELF executable for this

/// machine: a build for another target would not start, and the rollback

/// would only notice after three tries.

pub fn check_target(bytes: &[u8]) -> Result<()> {

    if bytes.len() < 20 || &bytes[..4] != b"\x7fELF" {

        bail!("not an ELF executable");

    }

    if bytes[4] != 2 || bytes[5] != 1 {

        bail!("not a 64-bit little-endian build");

    }

    let machine = u16::from_le_bytes([bytes[18], bytes[19]]);

    match machine_here() {

        Some(m) if m == machine => Ok(()),

        Some(_) => bail!("a build for another machine (ELF machine {machine}); this runner is {}", std::env::consts::ARCH),

        None => bail!("updates are not supported on {}", std::env::consts::ARCH),

    }

}

/// Run `file args…` with a time limit, returning its standard output.

fn run(file: &Path, args: &[&str], limit: Duration) -> Result<String> {

    let mut child = std::process::Command::new(file)

        .args(args)

        .stdin(std::process::Stdio::null())

        .stdout(std::process::Stdio::piped())

        .stderr(std::process::Stdio::piped())

        .spawn()

        .with_context(|| format!("starting {}", file.display()))?;

    let start = std::time::Instant::now();

    loop {

        if child.try_wait()?.is_some() {

            break;

        }

        if start.elapsed() > limit {

            let _ = child.kill();

            let _ = child.wait();

            bail!("`{} {}` did not finish within {} s", file.display(), args.join(" "), limit.as_secs());

        }

        std::thread::sleep(Duration::from_millis(50));

    }

    let out = child.wait_with_output()?;

    if !out.status.success() {

        bail!("`{} {}` failed ({}): {}", file.display(), args.join(" "), out.status, String::from_utf8_lossy(&out.stderr).trim());

    }

    Ok(String::from_utf8_lossy(&out.stdout).into_owned())

}

/// Check a build as the runner would run it: this machine's target, it

/// starts and names its version, and it finds the tiers this home offers.

/// Returns its version.

pub fn check(file: &Path, home: &Path) -> Result<String> {

    let mut head = vec![0u8; 64];

    let n = std::io::Read::read(&mut std::fs::File::open(file).with_context(|| format!("opening {}", file.display()))?, &mut head)?;

    check_target(&head[..n])?;

    let limit = Duration::from_secs(30);

    let line = run(file, &["--version"], limit)?;

    let version = parse_version(&line).with_context(|| format!("{} does not say it is flow: {:?}", file.display(), line.trim()))?;

    let home_s = home.to_string_lossy();

    let tiers = run(file, &["--home", &home_s, "tiers"], limit)?;

    if !tiers.lines().any(|l| l.ends_with(": available")) {

        bail!("flow {version} finds no tier it can run here:\n{}", tiers.trim());

    }

    Ok(version)

}

fn now_ms() -> u64 {

    std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_millis() as u64).unwrap_or(0)

}

fn log(home: &Path, event: &str, from: &str, to: &str, why: &str) {

    use std::io::Write;

    let line = serde_json::json!({"at_ms": now_ms(), "event": event, "from": from, "to": to, "why": why});

    if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open(home.join(LOG)) {

        let _ = writeln!(f, "{line}");

    }

}

/// The update not yet proven, if any.

pub fn pending(home: &Path) -> Option<Pending> {

    serde_json::from_slice(&std::fs::read(home.join(PENDING)).ok()?).ok()

}

fn save(home: &Path, p: &Pending) -> Result<()> {

    let tmp = home.join(format!("{PENDING}.tmp"));

    std::fs::write(&tmp, serde_json::to_vec_pretty(p)?)?;

    std::fs::rename(&tmp, home.join(PENDING))?;

    Ok(())

}

/// Put `staged` in place of `exe`, keeping `exe` as [`PREV`] beside it.

/// `exe` is never missing: it is copied aside, then replaced in one rename.

pub fn swap(exe: &Path, staged: &Path) -> Result<()> {

    let prev = exe.with_file_name(PREV);

    let _ = std::fs::remove_file(&prev);

    if std::fs::hard_link(exe, &prev).is_err() {

        std::fs::copy(exe, &prev).with_context(|| format!("keeping {} as {}", exe.display(), prev.display()))?;

    }

    std::fs::rename(staged, exe).with_context(|| format!("putting {} in place of {}", staged.display(), exe.display()))?;

    Ok(())

}

/// Put [`PREV`] back in place of `exe`.

pub fn restore(exe: &Path) -> Result<()> {

    let prev = exe.with_file_name(PREV);

    if !prev.is_file() {

        bail!("{} is not there to go back to", prev.display());

    }

    std::fs::rename(&prev, exe).with_context(|| format!("putting {} back", prev.display()))?;

    Ok(())

}

/// Check `file`, stage it beside `exe`, and unless `dry_run` swap it in and

/// tell the running runner (if any) to leave within `notice`, finishing

/// what fits and handing back the rest, so its unit starts the new build.

pub fn install(home: &crate::home::Home, exe: &Path, file: &Path, notice: Duration, dry_run: bool, allow_older: bool) -> Result<Installed> {

    if home.config.runner.as_ref().and_then(|r| r.update.as_ref()).is_some_and(|u| !u.allow) {

        bail!("flow.toml says [runner.update] allow = false: this runner takes no updates");

    }

    if let Some(p) = pending(&home.root) {

        bail!("the update from {} to {} is not proven yet; wait for it (or for its rollback)", p.from, p.to);

    }

    let from = env!("CARGO_PKG_VERSION").to_string();

    let staged = exe.with_file_name(NEW);

    let _ = std::fs::remove_file(&staged);

    std::fs::copy(file, &staged).with_context(|| format!("staging {} as {}", file.display(), staged.display()))?;

    {

        use std::os::unix::fs::PermissionsExt;

        std::fs::set_permissions(&staged, std::fs::Permissions::from_mode(0o755))?;

    }

    let checked = check(&staged, &home.root).and_then(|to| {

        if older(&to, &from) && !allow_older {

            bail!("flow {to} is older than this runner's {from} (--allow-older to go back)");

        }

        Ok(to)

    });

    let to = match checked {

        Ok(v) => v,

        Err(e) => {

            let _ = std::fs::remove_file(&staged);

            return Err(e);

        }

    };

    if dry_run {

        let _ = std::fs::remove_file(&staged);

        return Ok(Installed { from, to, exe: exe.to_path_buf(), told: None });

    }

    swap(exe, &staged)?;

    let p =

        Pending { from: from.clone(), to: to.clone(), exe: exe.to_path_buf(), at_ms: now_ms(), notice_secs: notice.as_secs(), starts: 0 };

    save(&home.root, &p)?;

    log(&home.root, "installed", &from, &to, "flow runner update --from");

    let told = running(&home.root).filter(|pid| unsafe { libc::kill(*pid as libc::pid_t, libc::SIGUSR1) } == 0);

    Ok(Installed { from, to, exe: exe.to_path_buf(), told })

}

/// The running runner's process id, from its pid file, if that process is

/// still a `flow runner`: a stale file's pid may be another program's now,

/// and SIGUSR1 would end it. Builds before #85 write no pid file (and do

/// not listen for SIGUSR1), so they are never told.

fn running(home: &Path) -> Option<u32> {

    let pid: u32 = std::fs::read_to_string(home.join(PID)).ok()?.trim().parse().ok()?;

    let cmdline = std::fs::read(format!("/proc/{pid}/cmdline")).ok()?;

    is_runner(&cmdline).then_some(pid)

}

/// Whether a NUL-separated command line is `flow … runner` with no subcommand.

fn is_runner(cmdline: &[u8]) -> bool {

    let args: Vec<&[u8]> = cmdline.split(|b| *b == 0).filter(|a| !a.is_empty()).collect();

    let Some((argv0, rest)) = args.split_first() else { return false };

    let name = argv0.rsplit(|b| *b == b'/').next().unwrap_or(argv0);

    name.starts_with(b"flow") && rest.last() == Some(&&b"runner"[..])

}

/// Note this process as the home's runner, for `flow runner update` to tell.

pub fn note_runner(home: &Path) {

    let _ = std::fs::write(home.join(PID), std::process::id().to_string());

}

/// The runner stopped: forget its pid.

pub fn forget_runner(home: &Path) {

    if std::fs::read_to_string(home.join(PID)).is_ok_and(|p| p.trim() == std::process::id().to_string()) {

        let _ = std::fs::remove_file(home.join(PID));

    }

}

/// What a runner starting up should do about an update.

#[derive(Debug, PartialEq, Eq)]

pub enum Start {

    /// No update waits to be proven.

    Settled,

    /// The new build is on trial: prove it within [`PROVE_WITHIN`].

    Trial(Pending),

    /// It failed too often; the previous build is back and should be run.

    RolledBack(Pending),

}

/// Count this start of an unproven build; after [`MAX_STARTS`], put the

/// previous build back.

pub fn on_start(home: &Path) -> Result<Start> {

    let Some(mut p) = pending(home) else { return Ok(Start::Settled) };

    p.starts += 1;

    if p.starts > MAX_STARTS {

        roll_back(home, &p, &format!("{MAX_STARTS} starts without reaching the hub"))?;

        return Ok(Start::RolledBack(p));

    }

    save(home, &p)?;

    Ok(Start::Trial(p))

}

/// In the runner's home: what to tell the hub once it answers.

const NOTE: &str = "update.note";

/// Keep `note` to tell the hub (a proof, a rollback) once it answers.

fn note(home: &Path, note: &str) {

    let _ = std::fs::write(home.join(NOTE), note);

}

/// What there is to tell the hub, once.

pub fn take_note(home: &Path) -> Option<String> {

    let n = std::fs::read_to_string(home.join(NOTE)).ok()?;

    let _ = std::fs::remove_file(home.join(NOTE));

    Some(n).filter(|n| !n.trim().is_empty())

}

/// The new build reached the hub: it stays.

pub fn proven(home: &Path) {

    if let Some(p) = pending(home) {

        let _ = std::fs::remove_file(home.join(PENDING));

        log(home, "proven", &p.from, &p.to, "reached the hub");

        note(home, &format!("runs flow {} (was {}); proven", p.to, p.from));

        eprintln!("flow runner: flow {} proven (it reached the hub); flow {} stays as {PREV}", p.to, p.from);

    }

}

/// Put the previous build back and forget the update.

pub fn roll_back(home: &Path, p: &Pending, why: &str) -> Result<()> {

    restore(&p.exe)?;

    let _ = std::fs::remove_file(home.join(PENDING));

    log(home, "rolled back", &p.to, &p.from, why);

    note(home, &format!("rolled back from flow {} to {}: {why}", p.to, p.from));

    eprintln!("flow runner: rolled back from flow {} to {} ({why})", p.to, p.from);

    Ok(())

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn versions_are_read_and_compared() {

        assert_eq!(parse_version("flow 0.3.1 (local build)\n").as_deref(), Some("0.3.1"));

        assert_eq!(parse_version("flow 0.10.0"), Some("0.10.0".into()));

        assert_eq!(parse_version("oxy 0.6.0"), None);

        assert_eq!(parse_version("flow banana"), None);

        assert!(older("0.3.1", "0.10.0"), "by number, not by text");

        assert!(!older("0.3.1", "0.3.1"));

        assert!(older("0.3", "0.3.1"));

    }

    #[test]

    fn only_a_build_for_this_machine_passes() {

        let me = std::fs::read(std::env::current_exe().unwrap()).unwrap();

        check_target(&me[..64]).unwrap();

        assert!(check_target(b"#!/bin/sh\necho flow 9.9.9\n").unwrap_err().to_string().contains("not an ELF"));

        let mut other = me[..64].to_vec();

        other[18..20].copy_from_slice(&(if std::env::consts::ARCH == "riscv64" { 62u16 } else { 243u16 }).to_le_bytes());

        assert!(check_target(&other).unwrap_err().to_string().contains("another machine"));

    }

    #[test]

    fn a_swap_keeps_the_old_build_and_three_failed_starts_bring_it_back() {

        let dir = tempfile::tempdir().unwrap();

        let (bin, home) = (dir.path().join("bin"), dir.path().join("home"));

        std::fs::create_dir_all(&bin).unwrap();

        std::fs::create_dir_all(&home).unwrap();

        let exe = bin.join("flow");

        std::fs::write(&exe, "old").unwrap();

        std::fs::write(bin.join(NEW), "new").unwrap();

        swap(&exe, &bin.join(NEW)).unwrap();

        assert_eq!(std::fs::read_to_string(&exe).unwrap(), "new");

        assert_eq!(std::fs::read_to_string(bin.join(PREV)).unwrap(), "old");

        assert!(!bin.join(NEW).exists());

        assert_eq!(on_start(&home).unwrap(), Start::Settled, "no update, nothing to do");

        let p = Pending { from: "0.3.1".into(), to: "0.3.2".into(), exe: exe.clone(), at_ms: 1, notice_secs: 600, starts: 0 };

        save(&home, &p).unwrap();

        for n in 1..=MAX_STARTS {

            assert_eq!(on_start(&home).unwrap(), Start::Trial(Pending { starts: n, ..p.clone() }));

        }

        assert!(matches!(on_start(&home).unwrap(), Start::RolledBack(_)));

        assert_eq!(std::fs::read_to_string(&exe).unwrap(), "old", "the previous build is back");

        assert!(pending(&home).is_none());

        let log = std::fs::read_to_string(home.join(LOG)).unwrap();

        assert!(log.contains("rolled back") && log.contains("0.3.2"), "{log}");

    }

    #[test]

    fn only_a_running_runner_is_told() {

        assert!(is_runner(b"/usr/local/bin/flow\0runner\0"));

        assert!(is_runner(b"flow\0--home\0/srv/flow\0runner\0"));

        assert!(!is_runner(b"flow\0runner\0update\0--from\0x\0"), "not the update command itself");

        assert!(!is_runner(b"/usr/bin/vim\0runner\0"));

        assert!(!is_runner(b""));

        let dir = tempfile::tempdir().unwrap();

        std::fs::write(dir.path().join(PID), std::process::id().to_string()).unwrap();

        assert_eq!(running(dir.path()), None, "this test is not a flow runner");

    }

    #[test]

    fn a_proven_build_stays() {

        let dir = tempfile::tempdir().unwrap();

        let exe = dir.path().join("flow");

        std::fs::write(&exe, "new").unwrap();

        let p = Pending { from: "0.3.1".into(), to: "0.3.2".into(), exe: exe.clone(), at_ms: 1, notice_secs: 0, starts: 1 };

        save(dir.path(), &p).unwrap();

        proven(dir.path());

        assert!(pending(dir.path()).is_none());

        assert_eq!(on_start(dir.path()).unwrap(), Start::Settled);

        assert!(std::fs::read_to_string(dir.path().join(LOG)).unwrap().contains("proven"));

    }

}