//! 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";

/// In the runner's home: a new build is in place, leave for it. How a

/// runner is told on Windows, which has no SIGUSR1.

const REQUEST: &str = "update.request";

/// 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)

}

/// `exe`'s sibling named `stem`, with `exe`'s extension (`flow.prev.exe`

/// beside `flow.exe` on Windows, `flow.prev` beside `flow` elsewhere).

pub fn beside(exe: &Path, stem: &str) -> PathBuf {

    match exe.extension() {

        Some(e) => exe.with_file_name(format!("{stem}.{}", e.to_string_lossy())),

        None => exe.with_file_name(stem),

    }

}

/// The PE machine number of this build's architecture (Windows).

fn pe_machine_here() -> Option<u16> {

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

        "x86_64" => 0x8664,

        "aarch64" => 0xAA64,

        _ => return None,

    })

}

/// Refuse a file that is not a PE executable for this machine (Windows).

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

    if bytes.len() < 0x40 || &bytes[..2] != b"MZ" {

        bail!("not a Windows executable");

    }

    let at = u32::from_le_bytes([bytes[0x3c], bytes[0x3d], bytes[0x3e], bytes[0x3f]]) as usize;

    if bytes.len() < at + 6 || &bytes[at..at + 4] != b"PE\0\0" {

        bail!("not a Windows executable (no PE header in its first {} bytes)", bytes.len());

    }

    let machine = u16::from_le_bytes([bytes[at + 4], bytes[at + 5]]);

    match pe_machine_here() {

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

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

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

    }

}

/// 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())

}

/// Copy `from` to `to` if it is `size` bytes with BLAKE3 `blake3`, as the

/// hub attests a build: a part file renamed into place, so `to` is either

/// the whole checked build or absent.

pub fn copy_checked(from: &Path, to: &Path, size: u64, blake3: &str) -> Result<()> {

    let bytes = std::fs::read(from).with_context(|| format!("reading {}", from.display()))?;

    let got = blake3::hash(&bytes).to_hex().to_string();

    if bytes.len() as u64 != size || got != blake3 {

        bail!("{} is not the build the hub attests ({} bytes, BLAKE3 {}…)", from.display(), bytes.len(), &got[..12]);

    }

    let part = to.with_extension("partial");

    std::fs::write(&part, &bytes).with_context(|| format!("writing {}", part.display()))?;

    std::fs::rename(&part, to)?;

    Ok(())

}

/// 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; 4096];

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

    if cfg!(windows) {

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

    } else {

        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(())

}

/// Whether `p` is the program running now.

fn is_self(p: &Path) -> bool {

    let canon = |p: &Path| std::fs::canonicalize(p).ok();

    std::env::current_exe().ok().and_then(|me| canon(&me)).is_some_and(|me| Some(me) == canon(p))

}

/// Rename program `from` to `to`. On Windows a running program cannot

/// rename its own file (other programs can), so a copy of it, run as

/// `flow __rename`, does it then.

fn rename_exe(from: &Path, to: &Path) -> Result<()> {

    if cfg!(windows) && is_self(from) {

        let helper = beside(from, "flow.rename");

        std::fs::copy(from, &helper).with_context(|| format!("copying {} to rename it", from.display()))?;

        let out = std::process::Command::new(&helper).arg("__rename").arg(from).arg(to).output();

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

        let out = out.with_context(|| format!("starting {}", helper.display()))?;

        if !out.status.success() {

            bail!("{}", String::from_utf8_lossy(&out.stderr).trim());

        }

        return Ok(());

    }

    Ok(crate::portable::rename(from, to)?)

}

/// 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.

/// On Windows, where a running program cannot be replaced but can be

/// renamed, `exe` is renamed aside and `staged` into its place.

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

    let prev = beside(exe, PREV);

    if cfg!(windows) {

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

        rename_exe(exe, &prev).with_context(|| format!("keeping {} as {}", exe.display(), prev.display()))?;

        if let Err(e) = crate::portable::rename(staged, exe) {

            let _ = crate::portable::rename(&prev, exe);

            return Err(e).with_context(|| format!("putting {} in place of {}", staged.display(), exe.display()));

        }

        return Ok(());

    }

    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 = beside(exe, PREV);

    if !prev.is_file() {

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

    }

    if cfg!(windows) {

        // The failed build may be the one running: rename it aside first.

        let failed = beside(exe, "flow.failed");

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

        rename_exe(exe, &failed).with_context(|| format!("moving {} aside", exe.display()))?;

        return crate::portable::rename(&prev, exe).with_context(|| format!("putting {} back", 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 = beside(exe, NEW);

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

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

    crate::portable::set_mode(&staged, 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 });

    }

    if let Err(e) = swap(exe, &staged) {

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

        return Err(e);

    }

    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| cfg!(not(unix)) || tell(*pid));

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

}

/// Ask runner `pid` to leave for the new build (SIGUSR1).

#[cfg(unix)]

fn tell(pid: u32) -> bool {

    unsafe { libc::kill(pid as libc::pid_t, libc::SIGUSR1) == 0 }

}

/// On Windows: a request in the runner's home, which its loop reads.

#[cfg(not(unix))]

fn tell(_pid: u32) -> bool {

    false

}

/// In the runner's home: drain (claim nothing more, finish what it holds,

/// then stop), as SIGTERM asks on Linux: `flow runner drain`, for Windows.

const DRAIN: &str = "drain.request";

/// Ask the runner of `home` to drain; whether one is running to ask.

pub fn ask_to_drain(home: &Path) -> std::io::Result<bool> {

    std::fs::write(home.join(DRAIN), "")?;

    Ok(home.join(PID).is_file())

}

/// Whether the runner has been asked, by file, to drain (once).

pub fn drain_requested(home: &Path) -> bool {

    std::fs::remove_file(home.join(DRAIN)).is_ok()

}

/// Whether the runner has been asked, by file, to leave for a new build

/// (and forget the asking).

pub fn requested(home: &Path) -> bool {

    std::fs::remove_file(home.join(REQUEST)).is_ok()

}

/// 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.

#[cfg(unix)]

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)

}

/// On Windows: a runner of this home has a pid file, and is asked by a

/// request file it reads (a stale one is cleared when a runner starts).

#[cfg(not(unix))]

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

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

    std::fs::write(home.join(REQUEST), pid.to_string()).ok()?;

    Some(pid)

}

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

#[cfg_attr(not(unix), allow(dead_code))]

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::remove_file(home.join(REQUEST));

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

    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.