use std::collections::HashSet;
use std::fs;
use std::io;
use std::os::fd::OwnedFd;
use std::thread;
use std::time::{Duration, Instant};

use nix::sys::prctl;
use rustix::io::Errno;
use rustix::process::{self, Pid, PidfdFlags, Signal, WaitOptions};

const REAP_POLL_INTERVAL: Duration = Duration::from_millis(10);
const MAX_WALK_PASSES: usize = 16;

pub fn become_subreaper() -> io::Result<()> {
    Ok(prctl::set_child_subreaper(true)?)
}

pub fn terminate_and_reap_all(grace_period: Duration) -> io::Result<()> {
    signal_all_children(&[Signal::TERM, Signal::CONT]);

    if reap_until(Instant::now() + grace_period)? {
        return Ok(());
    }

    kill_and_reap_all()
}

// Not exact: a process that forks during the last pass can miss the signals.
// That is good enough, because every survivor receives SIGKILL after the grace period.
fn signal_all_children(signals: &[Signal]) {
    let mut signaled = HashSet::new();
    for _ in 0..MAX_WALK_PASSES {
        let root = process::getpid();
        if !signal_new_children(root, read_children(root), signals, &mut signaled) {
            return;
        }
    }
}

fn read_children(pid: Pid) -> Vec<Pid> {
    let Ok(tasks) = fs::read_dir(format!("/proc/{}/task", pid.as_raw_pid())) else {
        return Vec::new();
    };

    tasks
        .filter_map(|task| fs::read_to_string(task.ok()?.path().join("children")).ok())
        .flat_map(|children| parse_pids(&children))
        .collect()
}

fn parse_pids(text: &str) -> Vec<Pid> {
    text.split_whitespace()
        .filter_map(|n| Pid::from_raw(n.parse().ok()?))
        .collect()
}

fn signal_new_children(
    parent: Pid,
    children: Vec<Pid>,
    signals: &[Signal],
    signaled: &mut HashSet<Pid>,
) -> bool {
    let mut has_signaled_new = false;

    for child in children {
        if signaled.contains(&child) {
            continue;
        }
        let Some(pidfd) = open_if_child_of(parent, child) else {
            continue;
        };

        let children_before_signal = read_children(child);
        send_signals(&pidfd, signals);
        signaled.insert(child);
        has_signaled_new = true;

        signal_new_children(child, children_before_signal, signals, signaled);
    }

    has_signaled_new
}

fn open_if_child_of(parent: Pid, pid: Pid) -> Option<OwnedFd> {
    let pidfd = process::pidfd_open(pid, PidfdFlags::empty()).ok()?;
    (read_parent(pid) == Some(parent)).then_some(pidfd)
}

fn read_parent(pid: Pid) -> Option<Pid> {
    let stat = fs::read_to_string(format!("/proc/{}/stat", pid.as_raw_pid())).ok()?;
    let fields_after_comm = &stat[stat.rfind(')')? + 1..];
    let ppid = fields_after_comm.split_whitespace().nth(1)?;
    Pid::from_raw(ppid.parse().ok()?)
}

fn send_signals(pidfd: &OwnedFd, signals: &[Signal]) {
    for &signal in signals {
        let _ = process::pidfd_send_signal(pidfd, signal);
    }
}

fn reap_until(deadline: Instant) -> io::Result<bool> {
    loop {
        match process::wait(WaitOptions::NOHANG) {
            Ok(Some(_)) => {}
            Ok(None) if Instant::now() < deadline => thread::sleep(REAP_POLL_INTERVAL),
            Ok(None) => return Ok(false),
            Err(Errno::CHILD) => return Ok(true),
            Err(e) => return Err(e.into()),
        }
    }
}

fn kill_and_reap_all() -> io::Result<()> {
    loop {
        signal_all_children(&[Signal::KILL]);

        match process::wait(WaitOptions::empty()) {
            Ok(_) => {}
            Err(Errno::CHILD) => return Ok(()),
            Err(e) => return Err(e.into()),
        }
    }
}