regy 0.1.0

use std::{
    collections::BTreeMap,
    ffi::{OsStr, OsString},
    path::{Path, PathBuf},
    process::Stdio,
    time::Duration,
};

use async_trait::async_trait;
use tokio::{
    io::{AsyncRead, AsyncReadExt},
    process::Command,
    time::timeout,
};

use crate::domain::errors::{AgentError, AgentResult, ErrorCode};

const MAX_STREAM_BYTES: usize = 1024 * 1024;

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct CommandSpec {
    pub program: PathBuf,
    pub args: Vec<OsString>,
    pub env: BTreeMap<OsString, OsString>,
    pub cwd: Option<PathBuf>,
    pub timeout: Duration,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct CommandOutput {
    pub status: i32,
    pub stdout: Vec<u8>,
    pub stderr: Vec<u8>,
}

#[async_trait]
pub(crate) trait CommandRunner: Send + Sync {
    async fn output(&self, spec: &CommandSpec) -> AgentResult<CommandOutput>;
}

#[derive(Debug, Default, Clone, Copy)]
pub(crate) struct TokioCommandRunner;

#[async_trait]
impl CommandRunner for TokioCommandRunner {
    async fn output(&self, spec: &CommandSpec) -> AgentResult<CommandOutput> {
        let program = program_basename(&spec.program);
        let mut command = Command::new(&spec.program);
        command
            .args(&spec.args)
            .envs(&spec.env)
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .kill_on_drop(true);
        if let Some(cwd) = &spec.cwd {
            command.current_dir(cwd);
        }

        let mut child = command
            .spawn()
            .map_err(|_| command_error(format!("command {program} failed to start")))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| command_error(format!("command {program} stdout unavailable")))?;
        let stderr = child
            .stderr
            .take()
            .ok_or_else(|| command_error(format!("command {program} stderr unavailable")))?;

        let execution = async {
            let (status, stdout, stderr) =
                tokio::join!(child.wait(), read_bounded(stdout), read_bounded(stderr),);
            (status, stdout, stderr)
        };
        let (status, stdout, stderr) = match timeout(spec.timeout, execution).await {
            Ok(result) => result,
            Err(_) => {
                if child.kill().await.is_err() {
                    let _ = child.wait().await;
                }
                return Err(command_error(format!("command {program} timed out")));
            }
        };

        let status =
            status.map_err(|_| command_error(format!("command {program} execution failed")))?;
        let status = status.code().unwrap_or(-1);
        let stdout_bytes = stream_byte_count(&stdout);
        let stderr_bytes = stream_byte_count(&stderr);

        if status != 0 {
            return Err(command_error(redacted_diagnostic(
                &program,
                status,
                stdout_bytes,
                stderr_bytes,
            )));
        }
        if stdout.is_err() || stderr.is_err() {
            return Err(command_error(format!(
                "command {program} output unavailable (status: {status}, stdout bytes: {stdout_bytes}, stderr bytes: {stderr_bytes})"
            )));
        }

        Ok(CommandOutput {
            status,
            stdout: stdout.expect("stdout result checked above"),
            stderr: stderr.expect("stderr result checked above"),
        })
    }
}

#[cfg(test)]
pub(crate) fn redact_command_output(program: &str, output: &CommandOutput) -> String {
    redacted_diagnostic(
        &command_label_basename(program),
        output.status,
        output.stdout.len(),
        output.stderr.len(),
    )
}

async fn read_bounded(mut stream: impl AsyncRead + Unpin) -> Result<Vec<u8>, StreamReadError> {
    let mut output = Vec::new();
    let mut chunk = [0_u8; 8192];
    loop {
        if output.len() == MAX_STREAM_BYTES {
            let mut extra = [0_u8; 1];
            return match stream.read(&mut extra).await {
                Ok(0) => Ok(output),
                Ok(_) | Err(_) => Err(StreamReadError {
                    bytes: output.len(),
                }),
            };
        }

        let remaining = MAX_STREAM_BYTES - output.len();
        let read_size = remaining.min(chunk.len());
        match stream.read(&mut chunk[..read_size]).await {
            Ok(0) => return Ok(output),
            Ok(read) => output.extend_from_slice(&chunk[..read]),
            Err(_) => {
                return Err(StreamReadError {
                    bytes: output.len(),
                });
            }
        }
    }
}

#[derive(Debug)]
struct StreamReadError {
    bytes: usize,
}

fn stream_byte_count(result: &Result<Vec<u8>, StreamReadError>) -> usize {
    match result {
        Ok(output) => output.len(),
        Err(error) => error.bytes,
    }
}

fn redacted_diagnostic(
    program: &str,
    status: i32,
    stdout_bytes: usize,
    stderr_bytes: usize,
) -> String {
    format!(
        "command {program} exited (status: {status}, stdout bytes: {stdout_bytes}, stderr bytes: {stderr_bytes})"
    )
}

fn program_basename(program: &Path) -> String {
    program
        .file_name()
        .unwrap_or_else(|| OsStr::new("command"))
        .to_string_lossy()
        .into_owned()
}

#[cfg(test)]
fn command_label_basename(program: &str) -> String {
    if program.is_empty()
        || !program.bytes().all(|byte| {
            byte.is_ascii_alphanumeric() || matches!(byte, b'/' | b'.' | b'_' | b'+' | b'-')
        })
    {
        return String::from("command");
    }

    let Some(basename) = Path::new(program).file_name().and_then(OsStr::to_str) else {
        return String::from("command");
    };
    if basename.is_empty()
        || !basename
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'+' | b'-'))
    {
        return String::from("command");
    }

    basename.to_owned()
}

fn command_error(message: String) -> AgentError {
    AgentError::new(ErrorCode::InvalidMessage, message)
}