chimera-core 0.2.0

use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::process::ExitStatus;

use serde::Serialize;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::task::JoinHandle;

use crate::AgentError;

const MAX_STDERR_SIZE: usize = 64 * 1024; // 64 KB

#[derive(Debug, Clone, bon::Builder)]
pub struct SubprocessCommand {
    pub binary: PathBuf,
    #[builder(default)]
    pub args: Vec<String>,
    #[builder(default)]
    pub env: HashMap<String, String>,
    /// Inherited environment variables to remove from the subprocess.
    /// Applied after `env` overlay; only meaningful when `inherit_env` is true.
    #[builder(default)]
    pub env_remove: HashSet<String>,
    /// If false, the child gets ONLY `env`. If true (default), inherits
    /// the current process env and overlays `env`.
    #[builder(default = true)]
    pub inherit_env: bool,
    pub cwd: Option<PathBuf>,
}

pub struct Subprocess {
    child: Child,
    stdin: Option<BufWriter<ChildStdin>>,
    stdout: Option<ChildStdout>,
    stderr_task: JoinHandle<String>,
}

impl Subprocess {
    pub fn spawn(cmd: SubprocessCommand) -> Result<Self, AgentError> {
        let mut command = Command::new(&cmd.binary);
        command.args(&cmd.args);

        if cmd.inherit_env {
            command.envs(&cmd.env);
            for key in &cmd.env_remove {
                command.env_remove(key);
            }
        } else {
            command.env_clear().envs(&cmd.env);
        }

        if let Some(cwd) = &cmd.cwd {
            command.current_dir(cwd);
        }

        command
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped());
        command.kill_on_drop(true);

        let mut child = command
            .spawn()
            .map_err(|source| AgentError::SpawnFailed { source })?;

        let stdin = child.stdin.take().map(BufWriter::new);
        let stdout = child.stdout.take();

        let stderr = child.stderr.take().expect("stderr was piped");
        let stderr_task = tokio::spawn(async move {
            let mut buf = vec![0u8; MAX_STDERR_SIZE];
            let mut reader = BufReader::new(stderr);
            let mut total = 0;

            loop {
                match reader.read(&mut buf[total..]).await {
                    Ok(0) => break,
                    Ok(n) => {
                        total += n;
                        if total >= MAX_STDERR_SIZE {
                            break;
                        }
                    }
                    Err(_) => break,
                }
            }

            String::from_utf8_lossy(&buf[..total]).into_owned()
        });

        Ok(Self {
            child,
            stdin,
            stdout,
            stderr_task,
        })
    }

    pub fn take_stdin(&mut self) -> Option<BufWriter<ChildStdin>> {
        self.stdin.take()
    }

    pub fn take_stdout(&mut self) -> Option<ChildStdout> {
        self.stdout.take()
    }

    pub async fn write(&mut self, data: &str) -> Result<(), AgentError> {
        let stdin = self.stdin.as_mut().ok_or_else(|| AgentError::Other {
            message: "stdin already closed".into(),
            source: None,
        })?;
        stdin
            .write_all(data.as_bytes())
            .await
            .map_err(|e| AgentError::StdinWrite { source: e })?;
        stdin
            .flush()
            .await
            .map_err(|e| AgentError::StdinWrite { source: e })
    }

    pub async fn write_json(&mut self, value: &impl Serialize) -> Result<(), AgentError> {
        let json = serde_json::to_string(value).map_err(|e| AgentError::Other {
            message: "failed to serialize JSON".into(),
            source: Some(Box::new(e)),
        })?;
        self.write(&json).await?;
        self.write("\n").await
    }

    pub async fn close_stdin(&mut self) {
        self.stdin.take();
    }

    pub async fn kill(&mut self) -> Result<(), AgentError> {
        self.child.kill().await.map_err(|e| AgentError::Other {
            message: "failed to kill process".into(),
            source: Some(Box::new(e)),
        })
    }

    pub async fn wait(mut self) -> Result<SubprocessOutput, AgentError> {
        let status = self.child.wait().await.map_err(|e| AgentError::Other {
            message: "failed to wait for process".into(),
            source: Some(Box::new(e)),
        })?;
        let stderr = self.stderr_task.await.unwrap_or_default();
        Ok(SubprocessOutput { status, stderr })
    }
}

#[derive(Debug)]
pub struct SubprocessOutput {
    pub status: ExitStatus,
    pub stderr: String,
}

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

    #[test]
    fn stderr_cap_threshold() {
        assert_eq!(MAX_STDERR_SIZE, 64 * 1024);
    }
}