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

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

/// 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::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 {}", p.to, p.from, beside(&p.exe, PREV).display());
    }
}

/// 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 a_windows_build_is_checked_as_pe_and_named_beside_with_its_extension() {
        let mut pe = vec![0u8; 0x100];
        pe[..2].copy_from_slice(b"MZ");
        pe[0x3c..0x40].copy_from_slice(&0x80u32.to_le_bytes());
        pe[0x80..0x84].copy_from_slice(b"PE\0\0");
        pe[0x84..0x86].copy_from_slice(&0x8664u16.to_le_bytes());
        if std::env::consts::ARCH == "x86_64" {
            check_pe(&pe).unwrap();
        }
        pe[0x84..0x86].copy_from_slice(&0x01c4u16.to_le_bytes());
        assert!(check_pe(&pe).is_err(), "another machine");
        assert!(check_pe(b"\x7fELF and more bytes than enough for a header, really, quite a few more")
            .unwrap_err()
            .to_string()
            .contains("not a Windows"));
        assert_eq!(beside(Path::new("C:/flow/flow.exe"), PREV), Path::new("C:/flow/flow.prev.exe"));
        assert_eq!(beside(Path::new("/usr/local/bin/flow"), NEW), Path::new("/usr/local/bin/flow.new"));
    }

    #[test]
    fn a_request_file_asks_once() {
        let dir = tempfile::tempdir().unwrap();
        assert!(!requested(dir.path()));
        std::fs::write(dir.path().join(REQUEST), "1").unwrap();
        assert!(requested(dir.path()) && !requested(dir.path()), "read once");
        std::fs::write(dir.path().join(REQUEST), "1").unwrap();
        note_runner(dir.path());
        assert!(!requested(dir.path()), "a starting runner clears a stale request");
    }

    #[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"));
    }
}