//! Self-update: a green CI run of the forge's own repo replaces the running binary and restarts.
//!
//! Event-driven. `maybe_deploy` runs from the runner finish path when a job finishes. Nothing
//! polls. After the swap the process exits with `RESTART_EXIT_CODE`, and the systemd unit
//! (`Restart=on-failure`) starts the new binary.
//!
//! Runners update the same way. The claim route advertises the running binary's sha256 to a runner
//! of the same target, and `GET /api/v1/ci/runners/binary` serves the bytes.

use std::path::{Path, PathBuf};
use std::sync::OnceLock;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use forge_core::ci::{CiRun, RunState};
use forge_core::config::SelfUpdateConfig;
use futures_util::StreamExt;
use futures_util::stream::BoxStream;
use http::header::{CONTENT_LENGTH, CONTENT_TYPE, HeaderValue};
use http_body_util::StreamBody;
use sha2::{Digest, Sha256};
use tokio::io::AsyncWriteExt;
use tokio_util::io::ReaderStream;
use topcoat::router::{Body, HeaderMap, StatusCode, response::Response};

use crate::state::AppState;

/// Exit code after a swap. systemd `Restart=on-failure` restarts the unit on any non-zero exit.
pub const RESTART_EXIT_CODE: u8 = 75;
/// Response header carrying the sha256 of the running binary, lowercase hex.
pub const SHA_HEADER: &str = "Forge-Binary-Sha256";
/// Request header a runner sends with its `<arch>-<os>` target.
pub const TARGET_HEADER: &str = "Forge-Runner-Target";
/// How long the staged binary may take to answer `--version`.
const SANITY_TIMEOUT: Duration = Duration::from_secs(10);

/// Serialises deploys. A deploy that finds another one running is skipped.
static DEPLOY_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
/// Wakes the server so it stops after a swap.
static RESTART: tokio::sync::Notify = tokio::sync::Notify::const_new();
static RESTART_REQUESTED: AtomicBool = AtomicBool::new(false);

/// The running binary, hashed once at startup.
pub struct Running {
    /// The executable path, used to place the sibling `.prev`, `.bad` and staging files.
    pub path: PathBuf,
    pub sha256: String,
}

static RUNNING: OnceLock<Option<Running>> = OnceLock::new();

/// The running binary, hashed on first use. `None` if it cannot be read, which disables
/// advertising and serving, and deploys.
pub fn running() -> Option<&'static Running> {
    RUNNING
        .get_or_init(|| match load_running() {
            Ok(running) => {
                tracing::info!(sha256 = %running.sha256, "self-update: running binary hashed");
                Some(running)
            }
            Err(e) => {
                tracing::error!(error = %e, "self-update: cannot hash the running binary");
                None
            }
        })
        .as_ref()
}

fn load_running() -> std::io::Result<Running> {
    let path = std::env::current_exe()?;
    let sha256 = sha256_reader(std::fs::File::open(image_path())?)?;
    Ok(Running { path, sha256 })
}

/// The file that holds the running image. On Linux that is `/proc/self/exe`, which still names the
/// old inode after a rename over the executable. Elsewhere it is the executable path.
fn image_path() -> PathBuf {
    let proc_exe = PathBuf::from("/proc/self/exe");
    if proc_exe.exists() {
        proc_exe
    } else {
        std::env::current_exe().unwrap_or(proc_exe)
    }
}

/// Resolves once a swap has asked the server to stop.
pub async fn restart_requested() {
    RESTART.notified().await;
}

/// Whether a swap asked the process to exit with `RESTART_EXIT_CODE`.
pub fn restart_pending() -> bool {
    RESTART_REQUESTED.load(Ordering::SeqCst)
}

/// `<arch>-<os>` for this build, as a runner reports it in `Forge-Runner-Target`.
pub fn target() -> String {
    format!("{}-{}", std::env::consts::ARCH, std::env::consts::OS)
}

/// Whether a runner's requested target is this build's target.
fn target_matches(requested: Option<&str>) -> bool {
    requested == Some(target().as_str())
}

/// Whether a finished run should deploy. Pure, so the rules can be tested. A run deploys only when
/// it is a successful push to the configured branch of the configured repo, and it is still the
/// newest run on that branch. Pull request runs (including fork runs) never deploy.
pub fn should_deploy(
    cfg: &SelfUpdateConfig,
    repo_full_name: &str,
    run: &CiRun,
    newest_on_ref: Option<&str>,
) -> bool {
    cfg.enabled
        && matches!(run.state, RunState::Success)
        && run.trigger == "push"
        && run.pull_number.is_none()
        && repo_full_name.eq_ignore_ascii_case(&cfg.repo)
        && run.git_ref == format!("refs/heads/{}", cfg.branch)
        && newest_on_ref == Some(run.id.as_str())
}

/// Why a candidate binary must not be installed, or `None` to install it. `bad_list` is the text
/// of `<exe>.bad`: one hex sha256 per line, `#` starts a comment.
pub fn skip_reason(candidate_sha: &str, running_sha: &str, bad_list: &str) -> Option<&'static str> {
    if candidate_sha.eq_ignore_ascii_case(running_sha) {
        return Some("same as the running binary");
    }
    let listed = bad_list
        .lines()
        .map(str::trim)
        .filter(|line| !line.is_empty() && !line.starts_with('#'))
        .any(|line| line.eq_ignore_ascii_case(candidate_sha));
    listed.then_some("listed in the bad list")
}

/// `<name>.<suffix>` next to `exe`, for example `forge.prev`.
fn sibling(exe: &Path, suffix: &str) -> PathBuf {
    let name = exe
        .file_name()
        .map(|name| name.to_string_lossy().into_owned())
        .unwrap_or_default();
    exe.with_file_name(format!("{name}.{suffix}"))
}

/// Called when a job finishes. Starts a deploy in the background when the run qualifies.
pub async fn maybe_deploy(state: &AppState, run: &CiRun) {
    let cfg = &state.config.self_update;
    if !cfg.enabled || !matches!(run.state, RunState::Success) {
        return;
    }
    let repo = match state.store.repo_by_id(&run.repo_id).await {
        Ok(Some(repo)) => repo,
        Ok(None) => return,
        Err(e) => {
            tracing::error!(error = %e, "self-update: loading the repo failed");
            return;
        }
    };
    let full_name = format!("{}/{}", repo.owner_name, repo.name);
    let newest = match state
        .store
        .latest_run_on_ref(&run.repo_id, &run.git_ref)
        .await
    {
        Ok(newest) => newest,
        Err(e) => {
            tracing::error!(error = %e, "self-update: loading the newest run failed");
            return;
        }
    };
    if !should_deploy(cfg, &full_name, run, newest.as_ref().map(|r| r.id.as_str())) {
        return;
    }
    tracing::info!(run_id = %run.id, "self-update: green push run on the deploy branch");
    let state = state.clone();
    let run = run.clone();
    tokio::spawn(async move { deploy(state, run).await });
}

/// Installs the run's artifact over the running binary and asks the server to restart. Logs and
/// returns on any failure, leaving the running binary in place.
pub async fn deploy(state: AppState, run: CiRun) {
    let Ok(_guard) = DEPLOY_LOCK.try_lock() else {
        tracing::info!(run_id = %run.id, "self-update: a deploy is already running, skipping");
        return;
    };
    if let Err(e) = deploy_inner(&state, &run).await {
        tracing::error!(run_id = %run.id, error = %e, "self-update: deploy failed");
    }
}

async fn deploy_inner(state: &AppState, run: &CiRun) -> Result<(), String> {
    let cfg = &state.config.self_update;
    let running = running().ok_or("the running binary was not hashed at startup")?;

    let jobs = state
        .store
        .jobs_for_run(&run.id)
        .await
        .map_err(|e| e.to_string())?;
    let Some(job) = jobs.iter().find(|job| job.name == cfg.job) else {
        tracing::info!(run_id = %run.id, job = %cfg.job, "self-update: no such job in the run, skipping");
        return Ok(());
    };
    let artifact = state
        .store
        .artifact(&job.id, &cfg.artifact)
        .await
        .map_err(|e| e.to_string())?;
    let Some(artifact) = artifact else {
        tracing::info!(run_id = %run.id, artifact = %cfg.artifact, "self-update: no such artifact, skipping");
        return Ok(());
    };
    let read = state
        .blobs
        .read(&crate::ci::artifact_key(&artifact.id))
        .await
        .map_err(|e| e.to_string())?
        .ok_or("the artifact bytes are missing")?;

    // Stage next to the executable, so the final rename stays on one filesystem.
    let staged = sibling(&running.path, "new");
    let sha = match write_staged(read.stream, &staged).await {
        Ok(sha) => sha,
        Err(e) => {
            let _ = std::fs::remove_file(&staged);
            return Err(e);
        }
    };
    if !sha.eq_ignore_ascii_case(&artifact.sha256) {
        let _ = std::fs::remove_file(&staged);
        return Err(format!(
            "artifact sha256 {sha} does not match the recorded {}",
            artifact.sha256
        ));
    }
    let bad = std::fs::read_to_string(sibling(&running.path, "bad")).unwrap_or_default();
    if let Some(reason) = skip_reason(&sha, &running.sha256, &bad) {
        let _ = std::fs::remove_file(&staged);
        tracing::info!(sha256 = %sha, reason, "self-update: skipped");
        return Ok(());
    }
    if let Err(e) = sanity_check(&staged).await {
        let _ = std::fs::remove_file(&staged);
        return Err(e);
    }
    if let Err(e) = swap(&running.path, &staged) {
        let _ = std::fs::remove_file(&staged);
        return Err(e);
    }
    tracing::info!(
        old_sha256 = %running.sha256,
        new_sha256 = %sha,
        exe = %running.path.display(),
        "self-update: binary replaced, restarting"
    );
    RESTART_REQUESTED.store(true, Ordering::SeqCst);
    RESTART.notify_one();
    Ok(())
}

/// Streams the artifact to `path`, hashing as it goes. The file is executable and synced before
/// this returns. Returns the sha256 in lowercase hex.
async fn write_staged(
    mut stream: BoxStream<'static, std::io::Result<bytes::Bytes>>,
    path: &Path,
) -> Result<String, String> {
    use std::os::unix::fs::PermissionsExt;

    let mut file = tokio::fs::File::create(path)
        .await
        .map_err(|e| format!("cannot create {}: {e}", path.display()))?;
    let mut hasher = Sha256::new();
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|e| e.to_string())?;
        hasher.update(&chunk);
        file.write_all(&chunk).await.map_err(|e| e.to_string())?;
    }
    file.flush().await.map_err(|e| e.to_string())?;
    file.sync_all().await.map_err(|e| e.to_string())?;
    drop(file);
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o755))
        .map_err(|e| e.to_string())?;
    Ok(hex::encode(hasher.finalize()))
}

/// Runs `<staged> --version` with a timeout. Only exit 0 passes.
async fn sanity_check(path: &Path) -> Result<(), String> {
    use std::process::Stdio;

    let mut child = tokio::process::Command::new(path)
        .arg("--version")
        .stdin(Stdio::null())
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .kill_on_drop(true)
        .spawn()
        .map_err(|e| format!("the new binary does not run: {e}"))?;
    match tokio::time::timeout(SANITY_TIMEOUT, child.wait()).await {
        Ok(Ok(status)) if status.success() => Ok(()),
        Ok(Ok(status)) => Err(format!("the new binary exited with {status} for --version")),
        Ok(Err(e)) => Err(format!("waiting for the new binary failed: {e}")),
        Err(_) => Err("the new binary did not answer --version within 10 seconds".into()),
    }
}

/// Keeps the current binary as `<exe>.prev` (a hard link when possible), then renames `staged`
/// over `exe`.
fn swap(exe: &Path, staged: &Path) -> Result<(), String> {
    let prev = sibling(exe, "prev");
    if prev.exists() {
        std::fs::remove_file(&prev).map_err(|e| format!("cannot remove the old prev: {e}"))?;
    }
    if std::fs::hard_link(exe, &prev).is_err() {
        std::fs::copy(exe, &prev).map_err(|e| format!("cannot keep the previous binary: {e}"))?;
    }
    std::fs::rename(staged, exe).map_err(|e| format!("cannot replace the binary: {e}"))
}

/// Whether a claim response should advertise the binary: self-update is on and the runner asked
/// for this target. Adds `Forge-Binary-Sha256` to a 200 or 204 response.
pub fn advertise(state: &AppState, request: &HeaderMap, res: &mut Response) {
    if !state.config.self_update.enabled {
        return;
    }
    if !matches!(res.status(), StatusCode::OK | StatusCode::NO_CONTENT) {
        return;
    }
    let requested = request.get(TARGET_HEADER).and_then(|v| v.to_str().ok());
    if !target_matches(requested) {
        return;
    }
    let Some(running) = running() else {
        return;
    };
    if let Ok(value) = HeaderValue::from_str(&running.sha256) {
        res.headers_mut().insert(SHA_HEADER, value);
    }
}

/// The running binary as a streamed response, with its sha256. The caller checks auth and that
/// self-update is enabled. `Err` when the binary cannot be read.
pub async fn binary_response() -> Result<Response, String> {
    let running = running().ok_or("the running binary was not hashed at startup")?;
    let file = tokio::fs::File::open(image_path())
        .await
        .map_err(|e| e.to_string())?;
    let len = file.metadata().await.map_err(|e| e.to_string())?.len();
    let frames = ReaderStream::new(file).map(|chunk| chunk.map(http_body::Frame::data));
    let mut res = Response::new(Body::new(StreamBody::new(frames)));
    let headers = res.headers_mut();
    headers.insert(
        CONTENT_TYPE,
        HeaderValue::from_static("application/octet-stream"),
    );
    headers.insert(CONTENT_LENGTH, HeaderValue::from(len));
    if let Ok(value) = HeaderValue::from_str(&running.sha256) {
        headers.insert(SHA_HEADER, value);
    }
    Ok(res)
}

fn sha256_reader(mut reader: impl std::io::Read) -> std::io::Result<String> {
    let mut hasher = Sha256::new();
    let mut buf = vec![0u8; 64 * 1024];
    loop {
        let n = reader.read(&mut buf)?;
        if n == 0 {
            break;
        }
        hasher.update(&buf[..n]);
    }
    Ok(hex::encode(hasher.finalize()))
}

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

    fn cfg() -> SelfUpdateConfig {
        SelfUpdateConfig {
            enabled: true,
            repo: "Tahinli/forge".into(),
            branch: "main".into(),
            job: "build".into(),
            artifact: "dist_forge".into(),
        }
    }

    fn run(trigger: &str, git_ref: &str, state: RunState, pull_number: Option<i64>) -> CiRun {
        CiRun {
            id: "run-7".into(),
            repo_id: "repo-1".into(),
            number: 7,
            trigger: trigger.into(),
            git_ref: git_ref.into(),
            commit_oid: "a".repeat(40),
            pull_number,
            actor_id: "user-1".into(),
            state,
            config_error: None,
            created_at: "now".into(),
            finished_at: None,
        }
    }

    /// A push run of the deploy branch that is still the newest run.
    fn good() -> CiRun {
        run("push", "refs/heads/main", RunState::Success, None)
    }

    #[test]
    fn push_run_on_the_branch_of_the_repo_deploys() {
        assert!(should_deploy(
            &cfg(),
            "Tahinli/forge",
            &good(),
            Some("run-7")
        ));
    }

    #[test]
    fn pull_request_runs_never_deploy() {
        // A PR run carries a pull number and a refs/pull ref, and its trigger is not "push".
        let pr = run(
            "pull_request",
            "refs/pull/3/head",
            RunState::Success,
            Some(3),
        );
        assert!(!should_deploy(&cfg(), "Tahinli/forge", &pr, Some("run-7")));
        // Even if a PR run were stored on the branch ref, the trigger and pull number exclude it.
        let pr_on_branch = run(
            "pull_request",
            "refs/heads/main",
            RunState::Success,
            Some(3),
        );
        assert!(!should_deploy(
            &cfg(),
            "Tahinli/forge",
            &pr_on_branch,
            Some("run-7")
        ));
        let push_with_pull = run("push", "refs/heads/main", RunState::Success, Some(3));
        assert!(!should_deploy(
            &cfg(),
            "Tahinli/forge",
            &push_with_pull,
            Some("run-7")
        ));
    }

    #[test]
    fn other_branch_tag_and_manual_runs_do_not_deploy() {
        let feature = run("push", "refs/heads/feature", RunState::Success, None);
        assert!(!should_deploy(
            &cfg(),
            "Tahinli/forge",
            &feature,
            Some("run-7")
        ));
        let tag = run("tag", "refs/tags/v1", RunState::Success, None);
        assert!(!should_deploy(&cfg(), "Tahinli/forge", &tag, Some("run-7")));
        let manual = run("manual", "refs/heads/main", RunState::Success, None);
        assert!(!should_deploy(
            &cfg(),
            "Tahinli/forge",
            &manual,
            Some("run-7")
        ));
    }

    #[test]
    fn other_repo_does_not_deploy() {
        assert!(!should_deploy(
            &cfg(),
            "Tahinli/other",
            &good(),
            Some("run-7")
        ));
        // Case differences in the owner or name still match the configured repo.
        assert!(should_deploy(
            &cfg(),
            "tahinli/FORGE",
            &good(),
            Some("run-7")
        ));
    }

    #[test]
    fn only_a_successful_run_deploys() {
        for state in [
            RunState::Failure,
            RunState::Cancelled,
            RunState::Error,
            RunState::Running,
        ] {
            let r = run("push", "refs/heads/main", state, None);
            assert!(!should_deploy(&cfg(), "Tahinli/forge", &r, Some("run-7")));
        }
    }

    #[test]
    fn an_older_run_on_the_branch_does_not_deploy() {
        assert!(!should_deploy(
            &cfg(),
            "Tahinli/forge",
            &good(),
            Some("run-9")
        ));
        assert!(!should_deploy(&cfg(), "Tahinli/forge", &good(), None));
    }

    #[test]
    fn disabled_config_never_deploys() {
        let mut off = cfg();
        off.enabled = false;
        assert!(!should_deploy(
            &off,
            "Tahinli/forge",
            &good(),
            Some("run-7")
        ));
    }

    #[test]
    fn same_sha_as_running_is_skipped() {
        let sha = "ab12";
        assert_eq!(
            skip_reason(sha, sha, ""),
            Some("same as the running binary")
        );
        assert_eq!(
            skip_reason("AB12", sha, ""),
            Some("same as the running binary")
        );
        assert_eq!(skip_reason("cd34", sha, ""), None);
    }

    #[test]
    fn bad_list_entries_are_skipped() {
        let bad = "# rolled back 2026-10-01\n1111\n  2222  \n\n";
        assert_eq!(
            skip_reason("2222", "ffff", bad),
            Some("listed in the bad list")
        );
        assert_eq!(
            skip_reason("1111", "ffff", bad),
            Some("listed in the bad list")
        );
        // A comment line is not a sha.
        assert_eq!(skip_reason("rolled", "ffff", bad), None);
        assert_eq!(skip_reason("3333", "ffff", bad), None);
    }

    #[test]
    fn target_must_match_this_build() {
        assert!(target_matches(Some(&target())));
        assert!(!target_matches(None));
        assert!(!target_matches(Some("riscv64-plan9")));
    }

    #[test]
    fn sha256_of_known_bytes() {
        let digest = sha256_reader(&b"abc"[..]).unwrap();
        assert_eq!(
            digest,
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );
    }
}