//! One-shot JSON subprocess helper.
//!
//! mcptest spawns user-supplied helper processes in two places that share
//! the same plain contract: the transform step and the
//! context-aware lifecycle hooks. Both spawn a command once,
//! write one JSON value to its stdin, close stdin so the child sees EOF,
//! read stdout to EOF, and parse it as a single JSON value. Factoring that
//! core out here keeps the contract identical for both and gives each
//! caller one place to add its own phase/command framing on top.
//!
//! The contract is deliberately one shot, unlike the  plugin host
//! which keeps a stateful child speaking newline-delimited frames. A spawn
//! failure, a non-zero exit, a timeout, or unparseable stdout is a typed
//! [`SubprocessError`]; the caller wraps it in its own message that names
//! the phase and command. Output is never logged here (it may carry
//! secrets); callers log only the command they ran.

use std::process::Stdio;
use std::time::Duration;

use serde_json::Value;
use thiserror::Error;
use tokio::io::AsyncWriteExt;
use tokio::process::Command;
use tokio::time::timeout;

/// Typed failure from running a one-shot JSON subprocess.
///
/// Each variant maps to a distinct caller-facing message. The caller adds
/// the phase and command line, so this enum stays free of that framing and
/// is reusable by both the transform and hook call sites.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum SubprocessError {
    /// The command could not be spawned (binary missing, permission
    /// denied).
    #[error("failed to spawn: {source}")]
    Spawn {
        /// Underlying OS error.
        #[source]
        source: std::io::Error,
    },
    /// The child did not finish within the supplied budget.
    #[error("timed out after {timeout:?}")]
    Timeout {
        /// Budget the run exceeded.
        timeout: Duration,
    },
    /// The child exited non-zero. Carries the stderr tail for diagnosis
    /// (stderr is the author's own log, not the JSON payload).
    #[error("exited with {code}: {stderr}")]
    NonZeroExit {
        /// Exit code, or a signal description when there was no code.
        code: String,
        /// Trimmed stderr from the child for the operator.
        stderr: String,
    },
    /// The child's stdout was not a single parseable JSON value.
    #[error("emitted unparseable JSON: {source}")]
    BadOutput {
        /// Underlying parse error.
        #[source]
        source: serde_json::Error,
    },
    /// An I/O error serializing stdin or collecting the child's output.
    #[error("I/O error: {source}")]
    Io {
        /// Underlying error.
        #[source]
        source: std::io::Error,
    },
}

/// Run `program` with `args` once, feed it `stdin_json`, and parse its
/// stdout as one JSON value.
///
/// Spawns the command with `env` set on top of the inherited parent
/// environment, writes `stdin_json` to stdin, closes stdin, waits up to
/// `budget` for the child to exit, and parses stdout. An empty stdout
/// returns `Ok(None)` so a caller can treat "the helper said nothing" as a
/// no-op rather than a parse error. Any spawn, timeout, exit, or parse
/// failure is a typed [`SubprocessError`] the caller wraps with its own
/// context.
pub async fn run_json_subprocess(
    program: &str,
    args: &[String],
    stdin_json: &Value,
    env: &[(String, String)],
    budget: Duration,
) -> Result<Option<Value>, SubprocessError> {
    let mut cmd = Command::new(program);
    cmd.args(args)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .kill_on_drop(true);
    for (name, value) in env {
        cmd.env(name, value);
    }

    let mut child = cmd
        .spawn()
        .map_err(|source| SubprocessError::Spawn { source })?;

    write_stdin(&mut child, stdin_json).await?;
    let output = match timeout(budget, child.wait_with_output()).await {
        Ok(Ok(output)) => output,
        Ok(Err(source)) => return Err(SubprocessError::Io { source }),
        Err(_) => return Err(SubprocessError::Timeout { timeout: budget }),
    };

    if !output.status.success() {
        return Err(SubprocessError::NonZeroExit {
            code: output
                .status
                .code()
                .map(|c| c.to_string())
                .unwrap_or_else(|| "a signal".to_string()),
            stderr: String::from_utf8_lossy(&output.stderr).trim().to_string(),
        });
    }

    parse_stdout(&output.stdout)
}

/// Parse the child's stdout. Whitespace-only output is `Ok(None)` so a
/// helper that intentionally emits nothing is a no-op rather than an error.
fn parse_stdout(stdout: &[u8]) -> Result<Option<Value>, SubprocessError> {
    if stdout.iter().all(u8::is_ascii_whitespace) {
        return Ok(None);
    }
    serde_json::from_slice(stdout)
        .map(Some)
        .map_err(|source| SubprocessError::BadOutput { source })
}

/// Serialize `input` and write it to the child's stdin, then drop the pipe
/// so the child sees EOF. Taking stdin out of the child here means
/// [`Command::wait_with_output`] does not deadlock waiting on a writer that
/// never closes.
async fn write_stdin(
    child: &mut tokio::process::Child,
    input: &Value,
) -> Result<(), SubprocessError> {
    let payload =
        serde_json::to_vec(input).map_err(|source| SubprocessError::BadOutput { source })?;
    let mut stdin = child.stdin.take().ok_or_else(|| SubprocessError::Io {
        source: std::io::Error::new(std::io::ErrorKind::BrokenPipe, "stdin not piped"),
    })?;
    match stdin.write_all(&payload).await {
        Ok(()) => {}
        // A child that exits before reading its stdin (a script that ignores
        // input and exits non-zero) closes the read end, so our write races
        // into a broken pipe. That is not the meaningful failure: swallow it
        // and let the caller wait for the child so its exit status drives the
        // result (a `NonZeroExit`, not an `Io` error). Without this the result
        // is timing-dependent across platforms. Other write errors still surface.
        Err(source) if source.kind() == std::io::ErrorKind::BrokenPipe => {}
        Err(source) => return Err(SubprocessError::Io { source }),
    }
    // Dropping `stdin` closes the pipe, signalling EOF to the child.
    drop(stdin);
    Ok(())
}

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

    fn sh(script: &str) -> (String, Vec<String>) {
        (
            "/bin/sh".to_string(),
            vec!["-c".to_string(), script.to_string()],
        )
    }

    #[tokio::test]
    async fn echoes_a_json_value_back() {
        let (p, a) = sh(r#"cat >/dev/null; echo '{"ok":true}'"#);
        let out = run_json_subprocess(&p, &a, &json!({}), &[], Duration::from_secs(5))
            .await
            .expect("runs");
        assert_eq!(out, Some(json!({"ok": true})));
    }

    #[tokio::test]
    async fn forwards_stdin_to_the_child() {
        // The child reflects whatever it reads on stdin straight back out.
        let (p, a) = sh("cat");
        let out = run_json_subprocess(&p, &a, &json!({"x": 1}), &[], Duration::from_secs(5))
            .await
            .expect("runs");
        assert_eq!(out, Some(json!({"x": 1})));
    }

    #[tokio::test]
    async fn empty_stdout_is_none_not_an_error() {
        let (p, a) = sh("cat >/dev/null; true");
        let out = run_json_subprocess(&p, &a, &json!({}), &[], Duration::from_secs(5))
            .await
            .expect("runs");
        assert_eq!(out, None);
    }

    #[tokio::test]
    async fn env_pairs_reach_the_child() {
        let (p, a) = sh(r#"cat >/dev/null; printf '{"v":"%s"}' "$MY_VAR""#);
        let env = vec![("MY_VAR".to_string(), "hi".to_string())];
        let out = run_json_subprocess(&p, &a, &json!({}), &env, Duration::from_secs(5))
            .await
            .expect("runs");
        assert_eq!(out, Some(json!({"v": "hi"})));
    }

    #[tokio::test]
    async fn non_zero_exit_carries_stderr() {
        let (p, a) = sh(r#"echo 'boom' >&2; exit 7"#);
        let err = run_json_subprocess(&p, &a, &json!({}), &[], Duration::from_secs(5))
            .await
            .expect_err("non-zero fails");
        match err {
            SubprocessError::NonZeroExit { code, stderr } => {
                assert_eq!(code, "7");
                assert!(stderr.contains("boom"), "{stderr}");
            }
            other => panic!("expected NonZeroExit, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn unparseable_stdout_is_bad_output() {
        let (p, a) = sh(r#"cat >/dev/null; echo 'not json'"#);
        let err = run_json_subprocess(&p, &a, &json!({}), &[], Duration::from_secs(5))
            .await
            .expect_err("bad output fails");
        assert!(matches!(err, SubprocessError::BadOutput { .. }), "{err:?}");
    }

    #[tokio::test]
    async fn missing_binary_is_a_spawn_error() {
        let err = run_json_subprocess(
            "this-binary-does-not-exist-mcptest",
            &[],
            &json!({}),
            &[],
            Duration::from_secs(5),
        )
        .await
        .expect_err("missing binary fails");
        assert!(matches!(err, SubprocessError::Spawn { .. }), "{err:?}");
    }

    #[tokio::test]
    async fn a_hung_child_times_out() {
        let (p, a) = sh("sleep 5");
        let err = run_json_subprocess(&p, &a, &json!({}), &[], Duration::from_millis(150))
            .await
            .expect_err("hang times out");
        assert!(matches!(err, SubprocessError::Timeout { .. }), "{err:?}");
    }
}
