//! External watchdog for a prebuilt Rust test or a direct reference worker.
//! Output is test progress, not proof of GPU progress. No total-runtime limit.
#[cfg(target_os = "macos")]
#[path = "../src/process_resources.rs"]
mod process_resources;

#[cfg(target_os = "macos")]
mod supervisor {
    use super::process_resources::{ResourceUsage, resource_usage};
    use serde_json::json;
    use std::io::{self, Read, Write};
    use std::os::unix::process::CommandExt;
    use std::process::{Child, Command, Stdio};
    use std::sync::{Arc, Mutex};
    use std::thread;
    use std::time::{Duration, Instant};

    const SAMPLE_INTERVAL: Duration = Duration::from_millis(100);

    pub(super) struct Limits {
        pub(super) bytes: u64,
        pub(super) start: Duration,
        pub(super) idle: Duration,
    }

    struct Worker(Child);

    impl Drop for Worker {
        fn drop(&mut self) {
            if !matches!(self.0.try_wait(), Ok(Some(_))) {
                // This process group was created for this child alone. Also
                // clean up helpers on error; memory accounting covers the test
                // process, so this launcher is not a process-tree supervisor.
                unsafe { libc::kill(-(self.0.id() as i32), libc::SIGKILL) };
                let _ = self.0.kill();
                let _ = self.0.wait();
            }
        }
    }

    #[derive(Default)]
    struct Output {
        last: Option<Instant>,
        failed: bool,
    }

    // Worker-side probes are observations, not inference progress. Buffer JSON
    // across pipe reads so a split heartbeat cannot keep a stuck worker alive.
    // Plain output (including Rust test banners without a newline) still counts.
    fn useful_output(pending: &mut Vec<u8>, bytes: &[u8]) -> bool {
        pending.extend_from_slice(bytes);
        let mut progressed = false;
        while let Some(end) = pending.iter().position(|&b| b == b'\n') {
            let line = &pending[..end];
            let observation = serde_json::from_slice::<serde_json::Value>(line)
                .ok()
                .is_some_and(|value| {
                    matches!(
                        value["event"].as_str(),
                        Some(
                            "gpu_canary_sample"
                                | "gpu_canary_summary"
                                | "reference_canary_summary"
                                | "test_resource_sample"
                        )
                    )
                });
            progressed |= !observation && line.iter().any(|b| !b.is_ascii_whitespace());
            pending.drain(..=end);
        }
        if pending.len() > 64 * 1024
            || pending
                .iter()
                .find(|b| !b.is_ascii_whitespace())
                .is_some_and(|&b| b != b'{')
        {
            progressed = true;
            pending.clear();
        }
        progressed
    }

    fn forward(
        mut input: impl Read + Send + 'static,
        mut output: impl Write + Send + 'static,
        progress: Arc<Mutex<Output>>,
    ) -> thread::JoinHandle<()> {
        thread::spawn(move || {
            let mut bytes = [0; 4096];
            let mut pending = Vec::new();
            loop {
                match input.read(&mut bytes) {
                    Ok(0) => return,
                    Ok(n) => {
                        if output
                            .write_all(&bytes[..n])
                            .and_then(|()| output.flush())
                            .is_err()
                        {
                            progress.lock().unwrap().failed = true;
                            return;
                        }
                        if useful_output(&mut pending, &bytes[..n]) {
                            progress.lock().unwrap().last = Some(Instant::now());
                        }
                    }
                    Err(error) if error.kind() == io::ErrorKind::Interrupted => continue,
                    Err(_) => {
                        progress.lock().unwrap().failed = true;
                        return;
                    }
                }
            }
        })
    }

    pub(super) fn run(command: &mut Command, limits: Limits) -> Result<(), String> {
        if limits.bytes == 0 || limits.start.is_zero() || limits.idle.is_zero() {
            return Err("memory and progress limits must be positive".into());
        }
        // Fail before spawning if process accounting is unavailable.
        resource_usage(std::process::id()).ok_or("process accounting unavailable")?;
        let mut worker = Worker(
            command
                .process_group(0)
                .stdin(Stdio::null())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped())
                .spawn()
                .map_err(|e| format!("could not start test: {e}"))?,
        );
        let progress = Arc::new(Mutex::new(Output::default()));
        let readers = [
            forward(
                worker.0.stdout.take().unwrap(),
                io::stdout(),
                progress.clone(),
            ),
            forward(
                worker.0.stderr.take().unwrap(),
                io::stderr(),
                progress.clone(),
            ),
        ];
        let started = Instant::now();
        let mut peak = ResourceUsage::default();
        eprintln!(
            "{}",
            json!({"event":"test_supervisor_started", "pid":worker.0.id(),
            "limit_bytes":limits.bytes, "start_seconds":limits.start.as_secs_f64(),
            "idle_seconds":limits.idle.as_secs_f64(), "sample_ms":100})
        );
        let outcome = loop {
            if let Some(status) = worker.0.try_wait().map_err(|e| e.to_string())? {
                break if status.success() {
                    Ok(())
                } else {
                    Err(format!("test exited with {status}"))
                };
            }
            let now = Instant::now();
            let Some(sample) = resource_usage(worker.0.id()) else {
                // Exit can race the sample; missing telemetry on a live worker
                // is a failure, never permission to keep running unmonitored.
                if let Some(status) = worker.0.try_wait().map_err(|e| e.to_string())? {
                    break if status.success() {
                        Ok(())
                    } else {
                        Err(format!("test exited with {status}"))
                    };
                }
                break Err("process accounting lost".into());
            };
            peak.include(sample);
            eprintln!(
                "{}",
                json!({"event":"test_resource_sample",
                "elapsed_ms":started.elapsed().as_millis(), "usage":sample})
            );
            if sample.physical_bytes > limits.bytes || sample.peak_physical_bytes > limits.bytes {
                break Err("memory_limit".into());
            }
            let output = progress.lock().map_err(|e| e.to_string())?;
            if output.failed {
                break Err("output monitoring failed".into());
            }
            let (last, timeout, reason) = output
                .last
                .map_or((started, limits.start, "start_timeout"), |last| {
                    (last, limits.idle, "continuation_timeout")
                });
            if now.saturating_duration_since(last) >= timeout {
                break Err(reason.into());
            }
            drop(output);
            thread::sleep(SAMPLE_INTERVAL);
        };
        drop(worker); // Kill and reap on every failure, including output errors.
        for reader in readers {
            reader.join().map_err(|_| "output reader panicked")?;
        }
        eprintln!(
            "{}",
            json!({"event":"test_resource_summary",
            "elapsed_ms":started.elapsed().as_millis(), "peak":peak,
            "error":outcome.as_ref().err()})
        );
        if progress.lock().map_err(|e| e.to_string())?.failed {
            return Err("output monitoring failed".into());
        }
        outcome
    }

    pub(super) fn cli() -> Result<(), String> {
        let args = std::env::args_os().skip(1).collect::<Vec<_>>();
        if args.len() < 5 || (args[3] != "--command" && args.len() != 5) {
            return Err("usage: test-supervisor MEMORY_MIB START_SECONDS IDLE_SECONDS TEST_BINARY TEST_FILTER\nOr: test-supervisor MEMORY_MIB START_SECONDS IDLE_SECONDS --command EXECUTABLE [ARG ...]\nMonitors the direct worker only; do not pass cargo or a spawning shell.".into());
        }
        let number = |i: usize| {
            args[i]
                .to_str()
                .and_then(|s| s.parse::<u64>().ok())
                .filter(|n| *n > 0)
                .ok_or("limits must be positive integers")
        };
        let limits = Limits {
            bytes: number(0)?
                .checked_mul(1024 * 1024)
                .ok_or("memory limit overflow")?,
            start: Duration::from_secs(number(1)?),
            idle: Duration::from_secs(number(2)?),
        };
        let mut command = worker_command(&args[3..])?;
        run(&mut command, limits)
    }

    fn worker_command(args: &[std::ffi::OsString]) -> Result<Command, String> {
        if args.len() < 2 || args[1].is_empty() {
            return Err("worker executable or test filter must not be empty".into());
        }
        if args[0] == "--command" {
            let mut command = Command::new(&args[1]);
            command.args(&args[2..]);
            Ok(command)
        } else if args.len() == 2 {
            let mut command = Command::new(&args[0]);
            command
                .arg(&args[1])
                .args(["--include-ignored", "--test-threads=1", "--nocapture"]);
            Ok(command)
        } else {
            Err("unexpected test arguments".into())
        }
    }

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

        #[test]
        fn direct_worker_arguments_are_not_rust_test_arguments() {
            let args = |a: &[&str]| a.iter().map(std::ffi::OsString::from).collect::<Vec<_>>();
            let direct = worker_command(&args(&[
                "--command",
                "/usr/bin/python3",
                "-u",
                "reference.py",
            ]))
            .unwrap();
            assert_eq!(direct.get_program(), "/usr/bin/python3");
            assert_eq!(
                direct.get_args().collect::<Vec<_>>(),
                ["-u", "reference.py"]
            );
            let test = worker_command(&args(&["test-binary", "fixture"])).unwrap();
            assert_eq!(
                test.get_args().collect::<Vec<_>>(),
                [
                    "fixture",
                    "--include-ignored",
                    "--test-threads=1",
                    "--nocapture"
                ]
            );
            assert!(worker_command(&args(&["--command", ""])).is_err());
            assert!(worker_command(&args(&["test-binary", "fixture", "extra"])).is_err());
        }

        #[test]
        fn watchdog_stops_memory_and_silence_but_not_total_runtime() {
            let limits = || Limits {
                bytes: 256 * 1024 * 1024,
                start: Duration::from_millis(300),
                idle: Duration::from_millis(300),
            };
            let mut sleeping = Command::new("/bin/sleep");
            sleeping.arg("5");
            assert_eq!(run(&mut sleeping, limits()).unwrap_err(), "start_timeout");
            let mut started = Command::new("/bin/sh");
            started.args(["-c", "printf ready; exec /bin/sleep 5"]);
            assert_eq!(
                run(&mut started, limits()).unwrap_err(),
                "continuation_timeout"
            );
            for (prefix, expected) in [
                ("", "start_timeout"),
                ("printf 'ready\\n'; ", "continuation_timeout"),
            ] {
                let mut heartbeat = Command::new("/bin/sh");
                heartbeat.args(["-c", &format!("{prefix}while :; do printf '{{\"event\":\"gpu_canary_sample\",\"ok\":true}}\\n'; sleep 0.05; done")]);
                assert_eq!(run(&mut heartbeat, limits()).unwrap_err(), expected);
            }
            let mut bounded = limits();
            bounded.bytes = 1;
            assert_eq!(run(&mut sleeping, bounded).unwrap_err(), "memory_limit");
            let mut progressing = Command::new("/bin/sh");
            progressing.args([
                "-c",
                "for i in 1 2 3 4 5 6 7 8; do printf progress; sleep 0.1; done",
            ]);
            let before = Instant::now();
            run(&mut progressing, limits()).unwrap();
            assert!(before.elapsed() >= Duration::from_millis(600));
            let mut failure = Command::new("/bin/sh");
            failure.args(["-c", "exit 7"]);
            assert!(
                run(&mut failure, limits())
                    .unwrap_err()
                    .contains("exit status: 7")
            );
        }

        #[test]
        fn split_probe_records_do_not_reset_progress() {
            let mut pending = Vec::new();
            assert!(!useful_output(&mut pending, b"{\"event\":\"gpu_canary_"));
            assert!(!useful_output(&mut pending, b"sample\",\"ok\":true}\n"));
            assert!(useful_output(
                &mut pending,
                b"{\"event\":\"mtp_tokens\",\"ids\":[42]}\n"
            ));
            assert!(useful_output(&mut pending, b"progress"));
            assert!(pending.is_empty());
        }
    }
}

fn main() {
    #[cfg(target_os = "macos")]
    if let Err(error) = supervisor::cli() {
        eprintln!("Test supervisor: {error}");
        std::process::exit(1);
    }
    #[cfg(not(target_os = "macos"))]
    {
        eprintln!("test-supervisor requires macOS process accounting");
        std::process::exit(1);
    }
}