use std::fs;
use std::io;
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);

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

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

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

    kill_and_reap_all()
}

fn signal_tree(parent: Pid, signals: &[Signal]) {
    for child in read_children(parent) {
        if signal_if_child_of(parent, child, signals) {
            signal_tree(child, signals);
        }
    }
}

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_if_child_of(parent: Pid, pid: Pid, signals: &[Signal]) -> bool {
    let Ok(pidfd) = process::pidfd_open(pid, PidfdFlags::empty()) else {
        return false;
    };

    read_parent(pid) == Some(parent)
        && signals
            .iter()
            .all(|&signal| process::pidfd_send_signal(&pidfd, signal).is_ok())
}

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 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_tree(process::getpid(), &[Signal::KILL]);

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