magi-code 0.96.1

use crate::cancellation::AgentCancellation;
use std::{
    io::{self, Read},
    process::{Child, ExitStatus},
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicUsize, Ordering},
        mpsc,
    },
    thread::{self, JoinHandle},
    time::{Duration, Instant},
};

#[cfg(unix)]
use std::process::Command;

const CLEANUP_POLL_INTERVAL: Duration = Duration::from_millis(10);
#[cfg(unix)]
const PROCESS_TABLE_PROBE_INTERVAL: Duration = Duration::from_millis(100);
#[cfg(unix)]
const TERM_GRACE: Duration = Duration::from_millis(50);
const KILL_GRACE: Duration = Duration::from_millis(1000);
pub(crate) const PIPE_READER_JOIN_TIMEOUT: Duration = Duration::from_millis(500);
#[cfg(unix)]
const PIPE_READER_POLL_TIMEOUT_MS: i32 = 100;
const PIPE_READER_STOP_TIMEOUT: Duration = Duration::from_millis(500);
static DETACHED_PIPE_READERS: AtomicUsize = AtomicUsize::new(0);

#[derive(Debug, Clone)]
pub(crate) struct CleanupOutcome {
    pub(crate) status: Option<ExitStatus>,
    pub(crate) cleanup_warning: Option<String>,
}

impl CleanupOutcome {
    fn new(status: Option<ExitStatus>, warnings: Vec<&'static str>) -> Self {
        let cleanup_warning = (!warnings.is_empty())
            .then(|| format!("process cleanup incomplete: {}", warnings.join("; ")));
        Self {
            status,
            cleanup_warning,
        }
    }
}

pub(crate) fn terminate_child_tree_and_wait(child: &mut Child) -> anyhow::Result<CleanupOutcome> {
    #[cfg(unix)]
    {
        terminate_unix_process_group(child)
    }
    #[cfg(not(unix))]
    {
        terminate_direct_child(child)
    }
}

#[cfg(unix)]
fn terminate_unix_process_group(child: &mut Child) -> anyhow::Result<CleanupOutcome> {
    let pgid = child.id();
    let mut warnings = Vec::new();
    let mut observed_status = child.try_wait()?;

    match signal_process_group("-TERM", pgid) {
        SignalResult::NoSuchProcess => return Ok(CleanupOutcome::new(observed_status, warnings)),
        SignalResult::Signaled => {}
        SignalResult::Failed => warnings.push("TERM signal command failed"),
    }

    if poll_process_group_exit(child, pgid, TERM_GRACE, &mut observed_status, &mut warnings)? {
        return Ok(CleanupOutcome::new(observed_status, Vec::new()));
    }

    match signal_process_group("-KILL", pgid) {
        SignalResult::NoSuchProcess => return Ok(CleanupOutcome::new(observed_status, Vec::new())),
        SignalResult::Signaled => {}
        SignalResult::Failed => warnings.push("KILL signal command failed"),
    }

    let exited =
        poll_process_group_exit(child, pgid, KILL_GRACE, &mut observed_status, &mut warnings)?;
    if exited {
        return Ok(CleanupOutcome::new(observed_status, Vec::new()));
    }
    warnings.push("process group still present after cleanup grace period");
    Ok(CleanupOutcome::new(observed_status, warnings))
}

#[cfg(unix)]
fn poll_process_group_exit(
    child: &mut Child,
    pgid: u32,
    max_wait: Duration,
    observed_status: &mut Option<ExitStatus>,
    warnings: &mut Vec<&'static str>,
) -> anyhow::Result<bool> {
    let start = Instant::now();
    let mut last_process_table_probe = start;
    while start.elapsed() < max_wait {
        if observed_status.is_none() {
            *observed_status = child.try_wait()?;
        }
        match probe_process_group_with_process_table(
            pgid,
            last_process_table_probe.elapsed() >= PROCESS_TABLE_PROBE_INTERVAL,
        ) {
            SignalResult::NoSuchProcess => return Ok(true),
            SignalResult::Signaled => {
                if last_process_table_probe.elapsed() >= PROCESS_TABLE_PROBE_INTERVAL {
                    last_process_table_probe = Instant::now();
                }
            }
            SignalResult::Failed => warnings.push("process group exit probe command failed"),
        }
        thread::sleep(CLEANUP_POLL_INTERVAL);
    }
    if observed_status.is_none() {
        *observed_status = child.try_wait()?;
    }
    match probe_process_group(pgid) {
        SignalResult::NoSuchProcess => Ok(true),
        SignalResult::Signaled => Ok(false),
        SignalResult::Failed => {
            warnings.push("process group exit probe command failed");
            Ok(false)
        }
    }
}

#[cfg(unix)]
fn probe_process_group(pgid: u32) -> SignalResult {
    probe_process_group_with_process_table(pgid, true)
}

#[cfg(unix)]
fn probe_process_group_with_process_table(pgid: u32, include_process_table: bool) -> SignalResult {
    match signal_process_group("-0", pgid) {
        SignalResult::Signaled if include_process_table => {
            match process_group_has_live_members(pgid) {
                Some(false) => SignalResult::NoSuchProcess,
                Some(true) | None => SignalResult::Signaled,
            }
        }
        other => other,
    }
}

#[cfg(unix)]
fn process_group_has_live_members(pgid: u32) -> Option<bool> {
    let output = Command::new("/bin/ps")
        .args(["-o", "stat=", "-o", "pgid=", "-ax"])
        .output()
        .ok()?;
    if !output.status.success() {
        return None;
    }
    Some(process_statuses_have_live_member(
        &String::from_utf8_lossy(&output.stdout),
        pgid,
    ))
}

#[cfg(unix)]
fn process_statuses_have_live_member(ps_output: &str, pgid: u32) -> bool {
    ps_output.lines().any(|line| {
        let mut fields = line.split_whitespace();
        let Some(status) = fields.next() else {
            return false;
        };
        let Some(process_group) = fields.next() else {
            return false;
        };
        process_group.parse::<u32>().ok() == Some(pgid) && !status.starts_with('Z')
    })
}

#[cfg(unix)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SignalResult {
    Signaled,
    NoSuchProcess,
    Failed,
}

#[cfg(unix)]
fn signal_process_group(signal: &str, pgid: u32) -> SignalResult {
    let process_group = format!("-{pgid}");
    match Command::new("/bin/kill")
        .arg(signal)
        .arg("--")
        .arg(process_group)
        .output()
    {
        Ok(output) if output.status.success() => SignalResult::Signaled,
        Ok(output) if is_no_such_process(&output.stderr) => SignalResult::NoSuchProcess,
        Ok(_) | Err(_) => SignalResult::Failed,
    }
}

#[cfg(unix)]
fn is_no_such_process(stderr: &[u8]) -> bool {
    String::from_utf8_lossy(stderr)
        .to_ascii_lowercase()
        .contains("no such process")
}

// ponytail: non-Unix cancellation kills only the direct child, not descendants; Windows has no
// process-group equivalent without Job Objects, upgrade path: use winapi Job Objects with
// JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE if Windows becomes a supported target.
#[cfg(not(unix))]
fn terminate_direct_child(child: &mut Child) -> anyhow::Result<CleanupOutcome> {
    let mut warnings = Vec::new();
    let mut observed_status = child.try_wait()?;
    if observed_status.is_none() {
        if let Err(_error) = child.kill() {
            warnings.push("direct child kill command failed");
        }
    }
    let start = Instant::now();
    while observed_status.is_none() && start.elapsed() < KILL_GRACE {
        observed_status = child.try_wait()?;
        if observed_status.is_some() {
            break;
        }
        thread::sleep(CLEANUP_POLL_INTERVAL);
    }
    if observed_status.is_none() {
        warnings.push("direct child still present after cleanup grace period");
    }
    Ok(CleanupOutcome::new(observed_status, warnings))
}

#[derive(Debug, Clone, Copy)]
pub(crate) struct BoundedChildProcessLimits {
    pub(crate) stdout_max_bytes: usize,
    pub(crate) stderr_max_bytes: usize,
    pub(crate) timeout: Duration,
    pub(crate) poll_interval: Duration,
}

#[derive(Debug)]
pub(crate) struct BoundedChildProcessOutput {
    pub(crate) stdout: String,
    pub(crate) stderr: String,
    pub(crate) stdout_truncated: bool,
    pub(crate) stderr_truncated: bool,
    pub(crate) timed_out: bool,
    pub(crate) status: Option<ExitStatus>,
    pub(crate) cleanup_warning: Option<String>,
}

pub(crate) fn run_bounded_child_process(
    mut child: Child,
    limits: BoundedChildProcessLimits,
    cancellation: &AgentCancellation,
) -> anyhow::Result<BoundedChildProcessOutput> {
    let stdout_pipe = child
        .stdout
        .take()
        .ok_or_else(|| anyhow::anyhow!("child stdout pipe unavailable"))?;
    let stderr_pipe = child
        .stderr
        .take()
        .ok_or_else(|| anyhow::anyhow!("child stderr pipe unavailable"))?;
    let stdout_truncated = Arc::new(AtomicBool::new(false));
    let stdout_receiver = spawn_bounded_pipe_reader(
        stdout_pipe,
        limits.stdout_max_bytes,
        Arc::clone(&stdout_truncated),
    );
    let stderr_truncated = Arc::new(AtomicBool::new(false));
    let stderr_receiver = spawn_bounded_pipe_reader(
        stderr_pipe,
        limits.stderr_max_bytes,
        Arc::clone(&stderr_truncated),
    );

    let start = Instant::now();
    let mut timed_out = false;
    let cleanup = loop {
        if cancellation.is_canceled() {
            break terminate_child_tree_and_wait(&mut child)?;
        }
        if child.try_wait()?.is_some() {
            break terminate_child_tree_and_wait(&mut child)?;
        }
        if stdout_truncated.load(Ordering::Relaxed) || stderr_truncated.load(Ordering::Relaxed) {
            break terminate_child_tree_and_wait(&mut child)?;
        }
        if start.elapsed() >= limits.timeout {
            timed_out = true;
            break terminate_child_tree_and_wait(&mut child)?;
        }
        thread::sleep(limits.poll_interval);
    };

    let stdout_result = recv_pipe_reader_with_timeout(stdout_receiver);
    let stderr_result = recv_pipe_reader_with_timeout(stderr_receiver);
    let mut reader_warnings = Vec::new();
    if stdout_result.timed_out {
        reader_warnings.push("stdout reader still blocked after cleanup");
    }
    if stderr_result.timed_out {
        reader_warnings.push("stderr reader still blocked after cleanup");
    }
    if stdout_result.stopped_before_eof {
        reader_warnings.push("stdout reader stopped before pipe EOF after cleanup");
    }
    if stderr_result.stopped_before_eof {
        reader_warnings.push("stderr reader stopped before pipe EOF after cleanup");
    }
    if stdout_result.join_timed_out {
        reader_warnings.push("stdout pipe reader thread did not exit within grace period");
    }
    if stderr_result.join_timed_out {
        reader_warnings.push("stderr pipe reader thread did not exit within grace period");
    }
    let cleanup_warning =
        combine_cleanup_warning(cleanup.cleanup_warning.as_deref(), &reader_warnings);
    if cancellation.is_canceled() {
        cancellation.check()?;
    }

    Ok(BoundedChildProcessOutput {
        stdout: stdout_result.output,
        stderr: stderr_result.output,
        stdout_truncated: stdout_truncated.load(Ordering::Relaxed),
        stderr_truncated: stderr_truncated.load(Ordering::Relaxed),
        timed_out,
        status: cleanup.status,
        cleanup_warning,
    })
}

#[derive(Debug)]
struct PipeReadOutput {
    output: String,
    stopped_before_eof: bool,
}

fn read_bounded_pipe_until_stopped(
    mut pipe: impl Read,
    max_bytes: usize,
    truncated: Arc<AtomicBool>,
    stop: &AtomicBool,
    poll_fd: Option<i32>,
) -> PipeReadOutput {
    // Intentional bounded read: preserves truncation hardening; replace only with benchmarked faster behavior.
    let mut bytes = Vec::new();
    let mut buffer = [0u8; 8192];
    let mut stopped_before_eof = false;
    while bytes.len() < max_bytes {
        let remaining = max_bytes - bytes.len();
        let read_len = remaining.min(buffer.len());
        match pipe.read(&mut buffer[..read_len]) {
            Ok(0) => break,
            Ok(count) => bytes.extend_from_slice(&buffer[..count]),
            Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
            Err(error) if poll_fd.is_some() && error.kind() == io::ErrorKind::WouldBlock => {
                if !wait_for_pipe_retry(stop, poll_fd) {
                    stopped_before_eof = true;
                    break;
                }
            }
            Err(_) => break,
        }
    }
    if bytes.len() == max_bytes {
        let mut extra = [0u8; 1];
        loop {
            match pipe.read(&mut extra) {
                Ok(count) if count > 0 => {
                    truncated.store(true, Ordering::Relaxed);
                    break;
                }
                Ok(_) => break,
                Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                Err(error) if poll_fd.is_some() && error.kind() == io::ErrorKind::WouldBlock => {
                    if !wait_for_pipe_retry(stop, poll_fd) {
                        stopped_before_eof = true;
                        break;
                    }
                }
                Err(_) => break,
            }
        }
    }
    PipeReadOutput {
        output: String::from_utf8_lossy(&bytes).to_string(),
        stopped_before_eof,
    }
}

fn wait_for_pipe_retry(stop: &AtomicBool, poll_fd: Option<i32>) -> bool {
    if stop.load(Ordering::Acquire) {
        return false;
    }
    #[cfg(unix)]
    if let Some(fd) = poll_fd {
        let mut descriptor = libc::pollfd {
            fd,
            events: libc::POLLIN | libc::POLLHUP | libc::POLLERR,
            revents: 0,
        };
        loop {
            // SAFETY: descriptor points to one initialized pollfd for this call only.
            let result = unsafe { libc::poll(&mut descriptor, 1, PIPE_READER_POLL_TIMEOUT_MS) };
            if result >= 0 {
                break;
            }
            if io::Error::last_os_error().kind() != io::ErrorKind::Interrupted {
                thread::sleep(CLEANUP_POLL_INTERVAL);
                break;
            }
        }
    } else {
        thread::sleep(CLEANUP_POLL_INTERVAL);
    }
    #[cfg(not(unix))]
    let _ = poll_fd;
    !stop.load(Ordering::Acquire)
}

pub(crate) struct PipeReaderHandle {
    receiver: mpsc::Receiver<PipeReadOutput>,
    join_handle: JoinHandle<()>,
    stop: Arc<AtomicBool>,
}

pub(crate) struct PipeReaderResult {
    pub(crate) output: String,
    pub(crate) timed_out: bool,
    pub(crate) join_timed_out: bool,
    pub(crate) stopped_before_eof: bool,
}

#[cfg(unix)]
pub(crate) fn spawn_bounded_pipe_reader(
    pipe: impl Read + std::os::fd::AsRawFd + Send + 'static,
    max_bytes: usize,
    truncated: Arc<AtomicBool>,
) -> PipeReaderHandle {
    let fd = pipe.as_raw_fd();
    let poll_fd = set_nonblocking(fd).is_ok().then_some(fd);
    spawn_bounded_reader(pipe, max_bytes, truncated, poll_fd)
}

#[cfg(not(unix))]
pub(crate) fn spawn_bounded_pipe_reader(
    pipe: impl Read + Send + 'static,
    max_bytes: usize,
    truncated: Arc<AtomicBool>,
) -> PipeReaderHandle {
    spawn_bounded_reader(pipe, max_bytes, truncated, None)
}

fn spawn_bounded_reader(
    pipe: impl Read + Send + 'static,
    max_bytes: usize,
    truncated: Arc<AtomicBool>,
    poll_fd: Option<i32>,
) -> PipeReaderHandle {
    let (sender, receiver) = mpsc::sync_channel(1);
    let stop = Arc::new(AtomicBool::new(false));
    let reader_stop = Arc::clone(&stop);
    let join_handle = thread::spawn(move || {
        let output =
            read_bounded_pipe_until_stopped(pipe, max_bytes, truncated, &reader_stop, poll_fd);
        let _ = sender.send(output);
    });
    PipeReaderHandle {
        receiver,
        join_handle,
        stop,
    }
}

#[cfg(unix)]
pub(crate) fn set_nonblocking(fd: std::os::fd::RawFd) -> io::Result<()> {
    // SAFETY: `fd` belongs to a live child pipe, and fcntl does not retain the descriptor.
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
    if flags == -1 {
        return Err(io::Error::last_os_error());
    }
    // SAFETY: same live descriptor; F_SETFL updates only descriptor status flags.
    if unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) } == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

pub(crate) fn recv_thread_completion_with_timeout<T>(
    receiver: &mpsc::Receiver<T>,
    timeout: Duration,
) -> Option<T> {
    receiver.recv_timeout(timeout).ok()
}

pub(crate) fn recv_pipe_reader_with_timeout(handle: PipeReaderHandle) -> PipeReaderResult {
    match handle.receiver.recv_timeout(PIPE_READER_JOIN_TIMEOUT) {
        Ok(output) => completed_pipe_reader(handle, output),
        Err(mpsc::RecvTimeoutError::Disconnected) => {
            completed_pipe_reader(handle, empty_pipe_completion())
        }
        Err(mpsc::RecvTimeoutError::Timeout) => {
            handle.stop.store(true, Ordering::Release);
            match handle.receiver.recv_timeout(PIPE_READER_STOP_TIMEOUT) {
                Ok(output) => completed_pipe_reader(handle, output),
                Err(mpsc::RecvTimeoutError::Disconnected) => {
                    completed_pipe_reader(handle, empty_pipe_completion())
                }
                Err(mpsc::RecvTimeoutError::Timeout) => {
                    let detached_count = DETACHED_PIPE_READERS.fetch_add(1, Ordering::Relaxed) + 1;
                    eprintln!(
                        "magi-code warning: pipe reader thread did not exit after cleanup; detaching reader (detached pipe readers: {detached_count})"
                    );
                    // ponytail: generic/non-Unix Read implementations may remain blocking; production Unix child pipes are nonblocking and honor `stop`.
                    PipeReaderResult {
                        output: String::new(),
                        timed_out: true,
                        join_timed_out: true,
                        stopped_before_eof: false,
                    }
                }
            }
        }
    }
}

fn completed_pipe_reader(handle: PipeReaderHandle, completion: PipeReadOutput) -> PipeReaderResult {
    let _ = handle.join_handle.join();
    PipeReaderResult {
        output: completion.output,
        timed_out: false,
        join_timed_out: false,
        stopped_before_eof: completion.stopped_before_eof,
    }
}

fn empty_pipe_completion() -> PipeReadOutput {
    PipeReadOutput {
        output: String::new(),
        stopped_before_eof: false,
    }
}

pub(crate) fn combine_cleanup_warning(
    base: Option<&str>,
    extra_warnings: &[&'static str],
) -> Option<String> {
    let mut warnings = Vec::new();
    if let Some(base) = base.filter(|warning| !warning.is_empty()) {
        warnings.push(base.to_string());
    }
    warnings.extend(extra_warnings.iter().map(|warning| (*warning).to_string()));
    (!warnings.is_empty()).then(|| warnings.join("; "))
}