//! Self-update for an idle runner. Each claim response may advertise the sha256 of the forge's
//! binary. When that differs from ours, the runner stops claiming, lets its running jobs finish,
//! swaps its own binary for the forge's, and exits with [`RESTART_EXIT_CODE`] so systemd starts
//! the new one.
//!
//! [`UpdateState`] makes the decisions and does no I/O. [`install`] does the file work.

use std::io::Read;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;

use sha2::{Digest, Sha256};

use crate::client::{ClientError, ForgeClient};

/// The exit code after a successful swap. systemd restarts the runner on any non-zero code.
pub const RESTART_EXIT_CODE: u8 = 75;
/// How long the new binary gets to answer `--version`.
const SMOKE_TIMEOUT: Duration = Duration::from_secs(10);
/// The temp file, next to the running binary, that the download goes to.
const NEW_NAME: &str = ".forge-runner.new";
/// The previous binary, kept for a manual rollback.
const PREV_SUFFIX: &str = ".prev";
/// Optional list of bad binaries, one sha256 per line, next to the running binary.
const BAD_SUFFIX: &str = ".bad";

/// Decides whether an update is due. It remembers the last sha that failed, so the runner does
/// not retry it on every claim. A different advertised sha clears that memory.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UpdateState {
    enabled: bool,
    /// The sha256 of the running binary. Updates are disabled when it is unknown.
    current: Option<String>,
    pending: Option<String>,
    given_up: Option<String>,
}

impl UpdateState {
    /// `current` is the sha256 (lowercase hex) of the running binary.
    pub fn new(auto_update: bool, current: Option<String>) -> Self {
        Self {
            enabled: auto_update && current.is_some(),
            current,
            pending: None,
            given_up: None,
        }
    }

    /// Records the sha the forge advertised on a claim response. Returns true when this makes a
    /// new update due. An absent header changes nothing.
    pub fn observe(&mut self, advertised: Option<&str>) -> bool {
        if !self.enabled {
            return false;
        }
        let Some(advertised) = advertised else {
            return false;
        };
        if Some(advertised) == self.current.as_deref()
            || Some(advertised) == self.given_up.as_deref()
        {
            // Already running it, or it failed and the forge has not moved on.
            self.pending = None;
            return false;
        }
        if self.pending.as_deref() == Some(advertised) {
            return false;
        }
        self.pending = Some(advertised.to_owned());
        true
    }

    /// Whether updates are on at all. False when `auto_update` is off or the binary is unknown.
    pub fn enabled(&self) -> bool {
        self.enabled
    }

    /// The sha to update to, if an update is due.
    pub fn pending(&self) -> Option<&str> {
        self.pending.as_deref()
    }

    /// Stops trying `sha` until the forge advertises a different one.
    pub fn give_up(&mut self, sha: &str) {
        self.given_up = Some(sha.to_owned());
        self.pending = None;
    }
}

/// Whether a `.bad` list, one sha256 per line, names `sha`.
pub fn is_listed_bad(list: &str, sha: &str) -> bool {
    list.lines()
        .any(|line| line.trim().eq_ignore_ascii_case(sha))
}

/// The sha256 (lowercase hex) of the file at `path`.
pub fn file_sha256(path: &Path) -> std::io::Result<String> {
    let mut file = std::fs::File::open(path)?;
    let mut hasher = Sha256::new();
    let mut buf = vec![0u8; 64 * 1024];
    loop {
        let n = file.read(&mut buf)?;
        if n == 0 {
            break;
        }
        hasher.update(&buf[..n]);
    }
    Ok(hex::encode(hasher.finalize()))
}

/// The path of the running binary, and the sha256 of its contents.
pub fn locate_self() -> std::io::Result<(PathBuf, String)> {
    let exe = std::env::current_exe()?;
    let sha = file_sha256(&exe)?;
    Ok((exe, sha))
}

/// What [`install`] did when it did not fail.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Installed {
    /// The new binary is in place. The caller should exit with [`RESTART_EXIT_CODE`].
    Replaced,
    /// The `.bad` list names the advertised binary, so nothing was replaced.
    Listed,
}

/// Why an update failed. The current binary is untouched in every case.
#[derive(Debug, thiserror::Error)]
pub enum UpdateError {
    #[error("local io: {0}")]
    Io(#[from] std::io::Error),
    #[error("download: {0}")]
    Download(#[from] ClientError),
    #[error("sha256 mismatch: forge advertised {expected}, downloaded {actual}")]
    Mismatch { expected: String, actual: String },
    #[error("new binary failed its check: {0}")]
    Smoke(String),
}

/// Downloads the forge's binary next to `exe`, checks it, and swaps it in. The old binary is kept
/// as `<exe>.prev`. On any error the temp file is removed.
pub async fn install(
    client: &ForgeClient,
    runner_secret: &str,
    exe: &Path,
    expected: &str,
) -> Result<Installed, UpdateError> {
    let new = exe.with_file_name(NEW_NAME);
    // A temp file left by an earlier crash must not be reused.
    let _ = std::fs::remove_file(&new);
    let result = stage_and_swap(client, runner_secret, exe, &new, expected).await;
    if !matches!(result, Ok(Installed::Replaced)) {
        let _ = std::fs::remove_file(&new);
    }
    result
}

async fn stage_and_swap(
    client: &ForgeClient,
    runner_secret: &str,
    exe: &Path,
    new: &Path,
    expected: &str,
) -> Result<Installed, UpdateError> {
    let actual = client.download_binary(runner_secret, new).await?;
    if actual != expected {
        return Err(UpdateError::Mismatch {
            expected: expected.to_owned(),
            actual,
        });
    }
    std::fs::set_permissions(new, std::fs::Permissions::from_mode(0o755))?;
    std::fs::File::open(new)?.sync_all()?;
    smoke_test(new).await?;

    let bad = match std::fs::read_to_string(sibling(exe, BAD_SUFFIX)) {
        Ok(list) => list,
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => String::new(),
        Err(e) => return Err(e.into()),
    };
    if is_listed_bad(&bad, expected) {
        return Ok(Installed::Listed);
    }

    let prev = sibling(exe, PREV_SUFFIX);
    match std::fs::remove_file(&prev) {
        Ok(()) => {}
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
        Err(e) => return Err(e.into()),
    }
    // A hard link keeps the old inode alive at no cost. Copy only where links are not allowed.
    if std::fs::hard_link(exe, &prev).is_err() {
        std::fs::copy(exe, &prev)?;
    }
    std::fs::rename(new, exe)?;
    Ok(Installed::Replaced)
}

/// Runs `<bin> --version` and requires exit 0 within [`SMOKE_TIMEOUT`].
async fn smoke_test(bin: &Path) -> Result<(), UpdateError> {
    let output = tokio::time::timeout(
        SMOKE_TIMEOUT,
        tokio::process::Command::new(bin)
            .arg("--version")
            .stdin(Stdio::null())
            .kill_on_drop(true)
            .output(),
    )
    .await
    .map_err(|_| UpdateError::Smoke("--version did not finish in time".into()))??;
    if !output.status.success() {
        return Err(UpdateError::Smoke(format!(
            "--version exited with {}",
            output.status
        )));
    }
    Ok(())
}

/// `exe` with `suffix` appended to its file name, for example `forge.prev`.
fn sibling(exe: &Path, suffix: &str) -> PathBuf {
    let mut name = exe.as_os_str().to_owned();
    name.push(suffix);
    PathBuf::from(name)
}

#[cfg(test)]
mod tests {
    use super::*;

    const OURS: &str = "aaaa";
    const THEIRS: &str = "bbbb";

    fn state(auto_update: bool) -> UpdateState {
        UpdateState::new(auto_update, Some(OURS.to_owned()))
    }

    #[test]
    fn a_different_sha_becomes_pending_once() {
        let mut s = state(true);
        assert!(s.observe(Some(THEIRS)));
        assert_eq!(s.pending(), Some(THEIRS));
        // The same sha again is not news.
        assert!(!s.observe(Some(THEIRS)));
        assert_eq!(s.pending(), Some(THEIRS));
    }

    #[test]
    fn matching_sha_or_no_header_leaves_nothing_pending() {
        let mut s = state(true);
        assert!(!s.observe(Some(OURS)));
        assert_eq!(s.pending(), None);
        assert!(!s.observe(None));
        assert_eq!(s.pending(), None);
    }

    #[test]
    fn absent_header_keeps_a_pending_update() {
        let mut s = state(true);
        s.observe(Some(THEIRS));
        s.observe(None);
        assert_eq!(s.pending(), Some(THEIRS));
    }

    #[test]
    fn disabled_or_unknown_current_never_pends() {
        let mut off = state(false);
        assert!(!off.observe(Some(THEIRS)));
        assert_eq!(off.pending(), None);

        let mut unknown = UpdateState::new(true, None);
        assert!(!unknown.observe(Some(THEIRS)));
        assert_eq!(unknown.pending(), None);
    }

    #[test]
    fn a_failed_sha_is_not_retried_until_the_forge_moves_on() {
        let mut s = state(true);
        s.observe(Some(THEIRS));
        s.give_up(THEIRS);
        assert_eq!(s.pending(), None);

        // Claims keep advertising the same sha: nothing becomes pending.
        assert!(!s.observe(Some(THEIRS)));
        assert_eq!(s.pending(), None);

        // A different sha is due again.
        assert!(s.observe(Some("cccc")));
        assert_eq!(s.pending(), Some("cccc"));
    }

    #[test]
    fn going_back_to_the_failed_sha_cancels_a_pending_update() {
        let mut s = state(true);
        s.observe(Some(THEIRS));
        s.give_up(THEIRS);
        s.observe(Some("cccc"));
        s.observe(Some(THEIRS));
        assert_eq!(s.pending(), None);
    }

    #[test]
    fn bad_list_matches_whole_lines_case_insensitively() {
        let list = "# comment is not a sha\nBBBB\n  cccc  \n";
        assert!(is_listed_bad(list, THEIRS));
        assert!(is_listed_bad(list, "cccc"));
        assert!(!is_listed_bad(list, "bbb"));
        assert!(!is_listed_bad("", THEIRS));
    }

    #[test]
    fn sibling_appends_the_suffix() {
        assert_eq!(
            sibling(Path::new("/home/t/.local/bin/forge"), PREV_SUFFIX),
            PathBuf::from("/home/t/.local/bin/forge.prev")
        );
    }

    #[test]
    fn file_sha256_matches_a_known_digest() {
        let dir = tempfile::tempdir().expect("tempdir");
        let path = dir.path().join("f");
        std::fs::write(&path, b"abc").expect("write");
        assert_eq!(
            file_sha256(&path).expect("hash"),
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );
    }
}