//! The builtin executor through the real sandbox helper. The image source is a stub that returns
//! a rootfs built from host binaries, so no registry is needed. The tests skip, with a message,
//! when unprivileged user namespaces are unavailable.

use std::collections::BTreeMap;
use std::fs;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::{Duration, Instant};

use forge_ci::proto::{JobAssignment, Outcome};
use forge_runner::check::build_rootfs;
use forge_runner::config::Limits;
use forge_runner::executor::{BuiltinExecutor, ExecRequest, ExecResult, Executor, ImageSource};
use forge_runner::oci::{ImageError, PreparedImage, RegistryAuth};
use forge_runner::sandbox::job_cgroup_path;
use nix::sys::signal::{Signal, kill};
use nix::unistd::Pid;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

/// Hands out a fixed rootfs for every image, with an env that the job can observe.
struct LocalRootfs(PathBuf);

impl ImageSource for LocalRootfs {
    fn prepare<'a>(
        &'a self,
        _image: &'a str,
        _auth: Option<&'a RegistryAuth>,
        _log: &'a mpsc::Sender<Vec<u8>>,
    ) -> impl Future<Output = Result<PreparedImage, ImageError>> + Send + 'a {
        std::future::ready(Ok(PreparedImage {
            rootfs: self.0.clone(),
            digest: format!("sha256:{}", "0".repeat(64)),
            env: vec![
                "PATH=/usr/local/bin:/usr/bin:/bin".to_owned(),
                "FROM_IMAGE=yes".to_owned(),
            ],
            working_dir: None,
            user: None,
        }))
    }
}

fn userns_available() -> bool {
    let max = fs::read_to_string("/proc/sys/user/max_user_namespaces")
        .ok()
        .and_then(|s| s.trim().parse::<u64>().ok())
        .unwrap_or(0);
    if max == 0 {
        return false;
    }
    Command::new("unshare")
        .args(["--user", "--map-root-user", "true"])
        .status()
        .is_ok_and(|s| s.success())
}

/// Skips the test when user namespaces are unavailable. Returns `true` when the test should run.
fn can_run() -> bool {
    if userns_available() {
        true
    } else {
        eprintln!("skipping: unprivileged user namespaces are not available on this host");
        false
    }
}

fn job(run: &str) -> JobAssignment {
    JobAssignment {
        job_id: "job_01".to_owned(),
        job_token: "tok_01_secret_token".to_owned(),
        repo: "owner/name".to_owned(),
        clone_url: "https://forge.example/owner/name.git".to_owned(),
        commit: "a".repeat(40),
        git_ref: "refs/heads/main".to_owned(),
        job_name: "build".to_owned(),
        image: "docker.io/library/alpine:3".to_owned(),
        run: run.to_owned(),
        env: BTreeMap::from([("PLAIN".to_owned(), "visible".to_owned())]),
        secrets: BTreeMap::from([("MY_TOKEN".to_owned(), "hunter2-secret".to_owned())]),
        artifacts: Vec::new(),
        timeout_minutes: 5,
        network: false,
    }
}

/// An executor over a rootfs in `dir`, and a fresh workspace under `dir`.
fn executor(dir: &Path) -> (BuiltinExecutor<LocalRootfs>, PathBuf) {
    let rootfs = dir.join("rootfs");
    build_rootfs(&rootfs, &["sh", "sleep", "id", "touch"]).expect("build rootfs");
    let workspace = dir.join("workspace");
    fs::create_dir_all(&workspace).expect("create workspace");
    let helper = PathBuf::from(env!("CARGO_BIN_EXE_forge-runner"));
    let exec = BuiltinExecutor::new(
        LocalRootfs(rootfs),
        Limits::default(),
        helper,
        dir.join("jobs"),
    );
    (exec, workspace)
}

/// Runs `job` and returns the result with everything it logged.
async fn run_job(
    exec: &BuiltinExecutor<LocalRootfs>,
    job: &JobAssignment,
    workspace: &Path,
    cancel: CancellationToken,
) -> (ExecResult, String) {
    let (tx, mut rx) = mpsc::channel::<Vec<u8>>(256);
    let collector = tokio::spawn(async move {
        let mut all = Vec::new();
        while let Some(chunk) = rx.recv().await {
            all.extend_from_slice(&chunk);
        }
        all
    });
    let result = exec
        .run(ExecRequest {
            job,
            workspace,
            logs: tx,
            cancel,
        })
        .await;
    let logs = collector.await.expect("log collector");
    (result, String::from_utf8_lossy(&logs).into_owned())
}

#[tokio::test]
async fn job_runs_in_the_sandbox_and_its_output_comes_back() {
    let _serial = serial().await;
    if !can_run() {
        return;
    }
    let dir = tempfile::tempdir().expect("tempdir");
    let (exec, workspace) = executor(dir.path());
    let run = "echo hi-from-builtin; echo \"env=$FROM_IMAGE,$PLAIN,$CI\"; \
               echo \"secret=$MY_TOKEN\"; id -u; touch /workspace/out.txt";
    let (result, logs) = run_job(&exec, &job(run), &workspace, CancellationToken::new()).await;

    assert_eq!(result.outcome, Outcome::Success, "{result:?}\n{logs}");
    assert_eq!(result.exit_code, Some(0));
    assert!(logs.contains("hi-from-builtin"), "{logs}");
    assert!(logs.contains("env=yes,visible,true"), "{logs}");
    assert!(logs.contains("secret=hunter2-secret"), "{logs}");
    assert!(
        logs.lines().any(|l| l == "0"),
        "not root in the job: {logs}"
    );
    assert!(logs.contains("image sha256:"), "{logs}");
    assert!(
        workspace.join("out.txt").exists(),
        "workspace write missing"
    );
    assert!(!dir.path().join("jobs/job_01").exists(), "job dir was kept");
}

#[tokio::test]
async fn nonzero_exit_is_a_failure_with_its_code() {
    let _serial = serial().await;
    if !can_run() {
        return;
    }
    let dir = tempfile::tempdir().expect("tempdir");
    let (exec, workspace) = executor(dir.path());
    let (result, logs) = run_job(&exec, &job("exit 3"), &workspace, CancellationToken::new()).await;
    assert_eq!(result.outcome, Outcome::Failure, "{logs}");
    assert_eq!(result.exit_code, Some(3));
}

#[tokio::test]
async fn cancel_stops_the_container() {
    let _serial = serial().await;
    if !can_run() {
        return;
    }
    let dir = tempfile::tempdir().expect("tempdir");
    let (exec, workspace) = executor(dir.path());
    let cancel = CancellationToken::new();
    let stop = cancel.clone();
    tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(1500)).await;
        stop.cancel();
    });
    let started = Instant::now();
    let (result, logs) = run_job(&exec, &job("sleep 120"), &workspace, cancel).await;
    assert_eq!(result.outcome, Outcome::Cancelled, "{logs}");
    // SIGTERM makes the helper kill its container, so this does not wait out the grace period.
    assert!(
        started.elapsed() < Duration::from_secs(25),
        "{:?}",
        started.elapsed()
    );
}

#[tokio::test]
async fn sandbox_setup_failure_is_an_error_with_the_reason() {
    let _serial = serial().await;
    if !can_run() {
        return;
    }
    let dir = tempfile::tempdir().expect("tempdir");
    let (exec, _workspace) = executor(dir.path());
    let missing = dir.path().join("no-such-workspace");
    let (result, logs) = run_job(&exec, &job("true"), &missing, CancellationToken::new()).await;
    assert_eq!(result.outcome, Outcome::Error, "{logs}");
    assert_eq!(result.exit_code, Some(125));
    let message = result.message.unwrap_or_default();
    assert!(
        message.contains("workspace is not a directory"),
        "{message}"
    );
}

/// The builtin tests run in one process. Each takes this lock, so the helper-finding test below
/// never sees another test's helper as its own child.
static SERIAL: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());

async fn serial() -> tokio::sync::MutexGuard<'static, ()> {
    SERIAL.lock().await
}

/// The pids of this process's children that are sandbox helpers.
fn helper_children() -> Vec<u32> {
    let me = std::process::id();
    let Ok(entries) = fs::read_dir("/proc") else {
        return Vec::new();
    };
    entries
        .filter_map(|e| e.ok()?.file_name().to_str()?.parse::<u32>().ok())
        .filter(|&pid| {
            let stat = fs::read_to_string(format!("/proc/{pid}/stat")).unwrap_or_default();
            // The fields after the command name in parentheses start with the state, then ppid.
            let ppid = stat
                .rsplit_once(')')
                .and_then(|(_, rest)| rest.split_whitespace().nth(1))
                .and_then(|p| p.parse::<u32>().ok());
            let cmdline = fs::read(format!("/proc/{pid}/cmdline")).unwrap_or_default();
            ppid == Some(me) && cmdline.windows(9).any(|w| w == b"__sandbox".as_slice())
        })
        .collect()
}

/// The job cgroup directory of the helper with pid `helper`, for job `job_01`.
fn job_cgroup_dir(helper: u32) -> PathBuf {
    let rel = job_cgroup_path("job_01", helper).expect("job cgroup path");
    Path::new("/sys/fs/cgroup").join(rel.strip_prefix("/").unwrap_or(&rel))
}

/// The pids listed in a cgroup's `cgroup.procs`.
fn procs_in(dir: &Path) -> Vec<u32> {
    fs::read_to_string(dir.join("cgroup.procs"))
        .unwrap_or_default()
        .lines()
        .filter_map(|l| l.trim().parse().ok())
        .collect()
}

/// Whether `pid` has exited. A zombie has exited too, it only waits to be reaped.
fn process_gone(pid: u32) -> bool {
    match fs::read_to_string(format!("/proc/{pid}/stat")) {
        Err(_) => true,
        Ok(stat) => {
            stat.rsplit_once(')')
                .and_then(|(_, rest)| rest.split_whitespace().next())
                == Some("Z")
        }
    }
}

/// The helper is SIGKILLed while its job is still running, so it cannot tear the container down.
/// The executor must then kill the job's processes, remove its cgroup, and remove its job dir.
#[tokio::test]
async fn sigkilled_helper_leaves_no_process_behind() {
    let _serial = serial().await;
    if !can_run() {
        return;
    }
    let dir = tempfile::tempdir().expect("tempdir");
    let (exec, workspace) = executor(dir.path());
    // `trap '' TERM` makes the job ignore SIGTERM. Only the helper's SIGKILL path can stop it.
    let run = "trap '' TERM; sleep 300";
    let job_spec = job(run);

    let killer = async {
        let deadline = Instant::now() + Duration::from_secs(30);
        let (helper, cgroup, procs) = loop {
            assert!(Instant::now() < deadline, "the job never started");
            let found = helper_children().into_iter().find_map(|pid| {
                let cgroup = job_cgroup_dir(pid);
                let procs = procs_in(&cgroup);
                (!procs.is_empty()).then_some((pid, cgroup, procs))
            });
            if let Some(found) = found {
                break found;
            }
            tokio::time::sleep(Duration::from_millis(50)).await;
        };
        kill(Pid::from_raw(helper as i32), Signal::SIGKILL).expect("kill the helper");
        (cgroup, procs)
    };

    let (run_result, (cgroup, procs)) = tokio::join!(
        run_job(&exec, &job_spec, &workspace, CancellationToken::new()),
        killer
    );
    let (result, logs) = run_result;

    assert_ne!(result.outcome, Outcome::Success, "{logs}");
    assert!(
        !cgroup.exists(),
        "job cgroup {} is still there",
        cgroup.display()
    );
    for pid in &procs {
        assert!(process_gone(*pid), "job process {pid} survived the helper");
    }
    assert!(!dir.path().join("jobs/job_01").exists(), "job dir was kept");
}