//! Process execution utilities for the smolvm agent.
//!
//! This module provides common helpers for spawning and managing child processes,
//! including timeout handling and output capture.

use std::io::Read;
use std::process::Child;
use std::time::{Duration, Instant};

const CONTAINER_INIT_ARG: &str = "container-init";

/// Whether this agent invocation is the image-independent init process used by
/// a persistent workload container.
pub fn container_init_requested() -> bool {
    std::env::args_os().nth(1).as_deref() == Some(std::ffi::OsStr::new(CONTAINER_INIT_ARG))
}

/// Run a minimal PID-1 reaper for a persistent workload container.
///
/// Foreground execs can create background descendants. When timeout cleanup
/// kills that process tree, the descendants are reparented to container PID 1.
/// A passive keepalive such as `tail -f /dev/null` never waits for them and
/// therefore leaks zombies across execs. Block the relevant signals and use
/// `sigwait`, which closes the usual reap-before-sleep race without requiring
/// an async signal handler.
#[cfg(target_os = "linux")]
pub fn run_container_init() -> i32 {
    let mut signals = unsafe { std::mem::zeroed::<libc::sigset_t>() };
    unsafe {
        libc::sigemptyset(&mut signals);
        libc::sigaddset(&mut signals, libc::SIGCHLD);
        libc::sigaddset(&mut signals, libc::SIGTERM);
        libc::sigaddset(&mut signals, libc::SIGINT);
    }
    if unsafe { libc::pthread_sigmask(libc::SIG_BLOCK, &signals, std::ptr::null_mut()) } != 0 {
        return 1;
    }

    loop {
        loop {
            let mut status = 0;
            let pid = unsafe { libc::waitpid(-1, &mut status, libc::WNOHANG) };
            if pid <= 0 {
                break;
            }
        }

        let mut signal = 0;
        if unsafe { libc::sigwait(&signals, &mut signal) } != 0 {
            return 1;
        }
        if signal == libc::SIGTERM || signal == libc::SIGINT {
            return 0;
        }
    }
}

#[cfg(not(target_os = "linux"))]
pub fn run_container_init() -> i32 {
    1
}

/// Exit code used when a command is killed due to timeout.
pub const TIMEOUT_EXIT_CODE: i32 = 124;

/// Map a process exit status to a numeric exit code.
///
/// Normal exit → the process's exit code. Terminated by a signal (where
/// `ExitStatus::code()` is `None`) → `128 + signal`, matching shell
/// convention (SIGINT → 130, SIGTERM → 143), so callers and scripts can
/// distinguish a signal kill from a normal exit instead of seeing an
/// opaque 255.
pub fn exit_code_from_status(status: &std::process::ExitStatus) -> i32 {
    use std::os::unix::process::ExitStatusExt;
    if let Some(code) = status.code() {
        code
    } else if let Some(sig) = status.signal() {
        128 + sig
    } else {
        -1
    }
}

/// Per-stream output cap for non-interactive exec. Vec<u8> is base64-encoded
/// in JSON frames (4/3 expansion). Two streams at this cap must fit within the
/// 32 MB frame limit with room for JSON overhead:
///   11 MiB × 2 × 4/3 ≈ 29.3 MiB encoded + ~2.7 MiB JSON headroom.
pub const MAX_EXEC_OUTPUT: usize = 11 * 1024 * 1024;

/// Maximum time to wait for reader threads to finish after the child is killed.
/// Guards against pathological cases where an inherited fd keeps a pipe open.
const READER_JOIN_TIMEOUT: Duration = Duration::from_secs(5);

/// Captured output from a child process.
#[derive(Debug, Default)]
pub struct ChildOutput {
    pub stdout: Vec<u8>,
    pub stderr: Vec<u8>,
}

/// Result of waiting for a child process.
#[derive(Debug)]
pub enum WaitResult {
    /// Process completed with the given exit code.
    Completed { exit_code: i32, output: ChildOutput },
    /// Process was killed due to timeout.
    TimedOut {
        output: ChildOutput,
        timeout_ms: u64,
    },
    /// Process was killed because the requesting client disconnected.
    /// Used to free the accept loop when the client gives up mid-exec.
    ClientDisconnected { output: ChildOutput },
}

/// Check whether the peer on `fd` has closed the connection.
///
/// Uses `recv(MSG_PEEK | MSG_DONTWAIT)` which is more reliable than `poll()`
/// on vsock — vsock's poll implementation doesn't always propagate POLLHUP
/// when the peer closes, but a zero-length peek is the canonical way to
/// detect half-closed sockets.
///
/// Returns `true` if the peer has closed OR the socket is in an error state.
/// Returns `false` if the socket is still alive OR we can't determine (fail
/// open — a bogus fd shouldn't cause us to kill a healthy child).
#[cfg(target_os = "linux")]
pub fn is_peer_closed(fd: std::os::unix::io::RawFd) -> bool {
    if fd < 0 {
        return false;
    }
    let mut buf = [0u8; 1];
    // SAFETY: buf is a valid write target, MSG_PEEK doesn't consume data.
    let rc = unsafe {
        libc::recv(
            fd,
            buf.as_mut_ptr() as *mut libc::c_void,
            buf.len(),
            libc::MSG_PEEK | libc::MSG_DONTWAIT,
        )
    };
    if rc == 0 {
        // Peer performed orderly shutdown (FIN received).
        return true;
    }
    if rc < 0 {
        let errno = std::io::Error::last_os_error().raw_os_error().unwrap_or(0);
        // EAGAIN = no data but connection alive → peer still there. (EWOULDBLOCK
        // is the same value as EAGAIN on Linux.) Any other error (ECONNRESET,
        // ENOTCONN, EBADF, etc.) → peer gone.
        return !matches!(errno, libc::EAGAIN);
    }
    // rc > 0: there's data in the buffer — connection is alive.
    false
}

#[cfg(not(target_os = "linux"))]
pub fn is_peer_closed(_fd: std::os::unix::io::RawFd) -> bool {
    false
}

/// Try to wait for a child process, handling EINTR by retrying.
///
/// EINTR can occur when a signal is delivered during the wait syscall.
/// This is not a real error - we should just retry the wait.
fn try_wait_with_eintr(child: &mut Child) -> std::io::Result<Option<std::process::ExitStatus>> {
    loop {
        match child.try_wait() {
            Ok(status) => return Ok(status),
            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
                // EINTR - signal interrupted the syscall, retry
                continue;
            }
            Err(e) => return Err(e),
        }
    }
}

/// Wait for a child process, killing it if the timeout expires OR if the
/// requesting client disconnects (indicated by `client_fd`, which is polled
/// each iteration).
///
/// Stdout and stderr are drained concurrently in background threads to prevent
/// pipe deadlock: if the child writes more than the OS pipe buffer (~64KB),
/// it blocks on write() while the agent blocks waiting for exit — neither side
/// makes progress. The background threads consume pipe data continuously,
/// preventing backpressure from stalling the child.
///
/// The client-disconnect check is the short-term mitigation for BUG-12/20:
/// when the host-side exec client is SIGTERM'd or times out, the agent's
/// accept loop was left blocked on the still-running child. Now we kill the
/// child as soon as we detect the peer has closed the connection, freeing
/// the accept loop for the next request.
pub fn wait_with_timeout_cleanup_and_liveness<F>(
    child: &mut Child,
    timeout_ms: Option<u64>,
    client_fd: Option<std::os::unix::io::RawFd>,
    on_abort: F,
) -> std::io::Result<WaitResult>
where
    F: FnOnce(),
{
    use std::sync::mpsc;

    const CHUNK_SIZE: usize = 64 * 1024;

    // Drain stdout/stderr in background threads BEFORE waiting for exit.
    // Threads send chunks via channels so the parent accumulates data
    // incrementally. On timeout/disconnect, already-received chunks are
    // preserved even if the reader thread is still blocked on a pipe that
    // hasn't reached EOF (e.g., background process inherited stdio).
    let (stdout_tx, stdout_rx) = mpsc::channel::<Vec<u8>>();
    let (stderr_tx, stderr_rx) = mpsc::channel::<Vec<u8>>();

    let stdout_handle = child.stdout.take().and_then(|mut out| {
        std::thread::Builder::new()
            .name("crun-stdout".into())
            .spawn(move || {
                let mut total = 0usize;
                loop {
                    let mut chunk = vec![0u8; CHUNK_SIZE];
                    match out.read(&mut chunk) {
                        Ok(0) => break, // EOF
                        Ok(n) => {
                            total += n;
                            chunk.truncate(n);
                            if stdout_tx.send(chunk).is_err() {
                                break; // receiver dropped
                            }
                            if total >= MAX_EXEC_OUTPUT {
                                break; // cap reached
                            }
                        }
                        Err(_) => break,
                    }
                }
            })
            .ok()
    });

    let stderr_handle = child.stderr.take().and_then(|mut err| {
        std::thread::Builder::new()
            .name("crun-stderr".into())
            .spawn(move || {
                let mut total = 0usize;
                loop {
                    let mut chunk = vec![0u8; CHUNK_SIZE];
                    match err.read(&mut chunk) {
                        Ok(0) => break,
                        Ok(n) => {
                            total += n;
                            chunk.truncate(n);
                            if stderr_tx.send(chunk).is_err() {
                                break;
                            }
                            if total >= MAX_EXEC_OUTPUT {
                                break;
                            }
                        }
                        Err(_) => break,
                    }
                }
            })
            .ok()
    });

    // Accumulated output — grows as reader threads send chunks.
    let mut stdout_buf = Vec::new();
    let mut stderr_buf = Vec::new();

    let poll_interval = Duration::from_millis(10);
    let deadline = timeout_ms.map(|ms| Instant::now() + Duration::from_millis(ms));

    // Drain any available chunks from the channels into local buffers.
    let drain_channels = |stdout_rx: &mpsc::Receiver<Vec<u8>>,
                          stderr_rx: &mpsc::Receiver<Vec<u8>>,
                          stdout_buf: &mut Vec<u8>,
                          stderr_buf: &mut Vec<u8>| {
        for chunk in stdout_rx.try_iter() {
            stdout_buf.extend_from_slice(&chunk);
        }
        for chunk in stderr_rx.try_iter() {
            stderr_buf.extend_from_slice(&chunk);
        }
    };

    loop {
        // Drain available chunks each iteration so local buffers stay current.
        drain_channels(&stdout_rx, &stderr_rx, &mut stdout_buf, &mut stderr_buf);

        match try_wait_with_eintr(child) {
            Ok(Some(status)) => {
                // Child exited — give reader threads a bounded window to finish.
                // After the child dies, pipe write ends close and readers see EOF.
                // Use is_finished() on handles to detect completion without consuming
                // chunks (try_recv as a probe races and can drop data).
                let join_deadline = Instant::now() + READER_JOIN_TIMEOUT;
                while Instant::now() < join_deadline {
                    drain_channels(&stdout_rx, &stderr_rx, &mut stdout_buf, &mut stderr_buf);
                    let stdout_done = stdout_handle.as_ref().is_none_or(|h| h.is_finished());
                    let stderr_done = stderr_handle.as_ref().is_none_or(|h| h.is_finished());
                    if stdout_done && stderr_done {
                        break;
                    }
                    std::thread::sleep(Duration::from_millis(10));
                }
                // Final drain after threads are done (or timed out).
                drain_channels(&stdout_rx, &stderr_rx, &mut stdout_buf, &mut stderr_buf);
                let exit_code = exit_code_from_status(&status);
                return Ok(WaitResult::Completed {
                    exit_code,
                    output: ChildOutput {
                        stdout: stdout_buf,
                        stderr: stderr_buf,
                    },
                });
            }
            Ok(None) => {
                if let Some(fd) = client_fd {
                    if is_peer_closed(fd) {
                        on_abort();
                        let _ = child.kill();
                        let _ = child.wait();
                        drain_channels(&stdout_rx, &stderr_rx, &mut stdout_buf, &mut stderr_buf);
                        return Ok(WaitResult::ClientDisconnected {
                            output: ChildOutput {
                                stdout: stdout_buf,
                                stderr: stderr_buf,
                            },
                        });
                    }
                }

                if let Some(deadline) = deadline {
                    if Instant::now() >= deadline {
                        on_abort();
                        let _ = child.kill();
                        let _ = child.wait();
                        drain_channels(&stdout_rx, &stderr_rx, &mut stdout_buf, &mut stderr_buf);
                        return Ok(WaitResult::TimedOut {
                            output: ChildOutput {
                                stdout: stdout_buf,
                                stderr: stderr_buf,
                            },
                            timeout_ms: timeout_ms.unwrap_or(0),
                        });
                    }
                }

                std::thread::sleep(poll_interval);
            }
            Err(e) => return Err(e),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_timeout_exit_code_value() {
        // Matches the standard timeout command exit code
        assert_eq!(TIMEOUT_EXIT_CODE, 124);
    }

    #[test]
    fn test_child_output_default() {
        let output = ChildOutput::default();
        assert!(output.stdout.is_empty());
        assert!(output.stderr.is_empty());
    }
}
