use std::fs;
use std::os::unix::process::ExitStatusExt;
use std::path::Path;
use std::process;
use std::sync::atomic::{AtomicU32, Ordering};

use tempfile::TempDir;

use crate::common::*;

#[test]
fn exit_code_of_entrypoint_is_returned() {
    let status = Sandbox::new(&[])
        .args(["--", "sh", "-c", "exit 42"])
        .status()
        .expect("agent-sandbox binary failed to execute");

    assert_eq!(status.code(), Some(42));
}

#[test]
fn explicit_exit_130_remains_an_exit_code() {
    let status = Sandbox::new(&[])
        .args(["--", "sh", "-c", "exit 130"])
        .status()
        .expect("agent-sandbox binary failed to execute");

    assert_eq!(status.code(), Some(130));
}

#[test]
fn signal_death_of_entrypoint_becomes_exit_code() {
    let status = Sandbox::new(&[])
        .args(["--", "sh", "-c", "kill -TERM $$"])
        .status()
        .expect("agent-sandbox binary failed to execute");

    assert_eq!(status.code(), Some(128 + libc::SIGTERM));
}

#[test]
fn signal_death_of_bwrap_is_mirrored() {
    let status = Sandbox::new(&["--blacklist"])
        .args(["--", "sh", "-c", "kill -KILL $PPID; sleep 5"])
        .status()
        .expect("agent-sandbox binary failed to execute");

    assert_eq!(status.signal(), Some(libc::SIGKILL));
}

#[test]
fn detached_child_process_is_killed_after_entrypoint_exits_in_blacklist_mode() {
    assert_detached_child_process_is_killed("--blacklist", "");
}

#[test]
fn detached_child_process_is_killed_after_entrypoint_exits_in_whitelist_mode() {
    assert_detached_child_process_is_killed("--whitelist", "");
}

#[test]
fn detached_child_process_that_ignores_sigterm_is_killed() {
    assert_detached_child_process_is_killed("--blacklist", "trap \"\" TERM;");
}

#[test]
fn detached_child_process_receives_sigterm_before_sigkill() {
    let dir = TempDir::new().expect("failed to create temp dir");
    let flag = dir.path().join("terminated");
    let setup = format!("trap \"touch {}; exit\" TERM;", flag.display());

    let marker = create_unique_sleep_duration();
    run_with_detached_sleeper(
        &["--blacklist", "--rw", dir.path().to_str().unwrap()],
        &setup,
        &marker,
    );

    assert!(
        flag.exists(),
        "detached child process did not receive SIGTERM"
    );
}

#[test]
fn every_process_in_a_large_tree_receives_sigterm() {
    let dir = TempDir::new().expect("failed to create temp dir");
    let marker = create_unique_sleep_duration();
    let node_script = dir.path().join("node.sh");
    fs::write(&node_script, TREE_NODE_SCRIPT.replace("{marker}", &marker)).unwrap();
    let script = format!(
        "sh {node} root 0 &
         until [ $(ls {dir} | grep -c '^ready\\.') -eq {TREE_SIZE} ]; do sleep 0.01; done",
        node = node_script.display(),
        dir = dir.path().display(),
    );

    let status = Sandbox::new(&["--blacklist", "--rw", dir.path().to_str().unwrap()])
        .args(["--", "sh", "-c", &script])
        .status()
        .expect("agent-sandbox binary failed to execute");
    assert!(status.success());

    let terminated = count_files_with_prefix(dir.path(), "terminated.");
    let leftovers = find_processes_with_arg(&marker);
    for pid in &leftovers {
        unsafe { libc::kill(*pid, libc::SIGKILL) };
    }
    assert_eq!(terminated, TREE_SIZE, "processes that received SIGTERM");
    assert!(leftovers.is_empty(), "leftover processes: {leftovers:?}");
}

fn assert_detached_child_process_is_killed(mode: &str, setup: &str) {
    let marker = create_unique_sleep_duration();
    run_with_detached_sleeper(&[mode], setup, &marker);

    let leftovers = find_processes_with_arg(&marker);
    for pid in &leftovers {
        unsafe { libc::kill(*pid, libc::SIGKILL) };
    }
    assert!(leftovers.is_empty(), "leftover processes: {leftovers:?}");
}

fn create_unique_sleep_duration() -> String {
    static NEXT: AtomicU32 = AtomicU32::new(0);
    let sequence = NEXT.fetch_add(1, Ordering::Relaxed);
    format!("3600.{}{sequence}", process::id())
}

fn run_with_detached_sleeper(sandbox_args: &[&str], setup: &str, marker: &str) {
    let script = format!(
        "setsid -f sh -c '{setup} sleep {marker} & wait' </dev/null >/dev/null 2>&1
         until pgrep -fx 'sleep {marker}' >/dev/null; do sleep 0.01; done"
    );

    let status = Sandbox::new(sandbox_args)
        .env("AGENT_SANDBOX_CLEANUP_GRACE_PERIOD_MS", "1000")
        .args(["--", "sh", "-c", &script])
        .status()
        .expect("agent-sandbox binary failed to execute");
    assert!(status.success());
}

fn find_processes_with_arg(arg: &str) -> Vec<i32> {
    fs::read_dir("/proc")
        .unwrap()
        .filter_map(|entry| entry.ok()?.file_name().to_str()?.parse().ok())
        .filter(|pid: &i32| {
            fs::read(format!("/proc/{pid}/cmdline"))
                .is_ok_and(|cmdline| cmdline.split(|&b| b == 0).any(|a| a == arg.as_bytes()))
        })
        .collect()
}

const TREE_SIZE: usize = 1 + 3 + 9 + 27 + 81;

const TREE_NODE_SCRIPT: &str = r#"
name=$1 depth=$2 dir=$(dirname "$0")
trap 'touch "$dir/terminated.$name"; exit' TERM
if [ "$depth" -lt 4 ]; then
    setsid -f sh "$0" "$name.1" $((depth + 1)) </dev/null >/dev/null 2>&1
    sh "$0" "$name.2" $((depth + 1)) &
    sh "$0" "$name.3" $((depth + 1)) &
fi
sleep {marker} &
touch "$dir/ready.$name"
wait
"#;

fn count_files_with_prefix(dir: &Path, prefix: &str) -> usize {
    fs::read_dir(dir)
        .unwrap()
        .filter(|entry| {
            entry
                .as_ref()
                .is_ok_and(|e| e.file_name().to_string_lossy().starts_with(prefix))
        })
        .count()
}