eli 0.3.2

//! Managed shell processes for the `bash` tool.

use std::collections::HashMap;
use std::sync::Arc;

use tokio::io::AsyncReadExt;
use tokio::process::Command;
use tokio::sync::Mutex;

// ---------------------------------------------------------------------------
// ManagedShell
// ---------------------------------------------------------------------------

/// A shell process with captured stdout/stderr.
#[derive(Debug)]
pub struct ManagedShell {
    pub shell_id: String,
    pub cmd: String,
    pub cwd: Option<String>,
    pub output_chunks: Vec<String>,
    child: Option<tokio::process::Child>,
    exit_code: Option<i32>,
    read_handles: Vec<tokio::task::JoinHandle<Vec<String>>>,
}

impl ManagedShell {
    /// Concatenated output of stdout and stderr.
    pub fn output(&self) -> String {
        self.output_chunks.join("")
    }

    /// Process return code, or `None` if still running.
    pub fn returncode(&self) -> Option<i32> {
        self.exit_code
    }

    /// Human-readable status.
    pub fn status(&self) -> &'static str {
        if self.exit_code.is_some() {
            "exited"
        } else {
            "running"
        }
    }
}

// ---------------------------------------------------------------------------
// ShellManager
// ---------------------------------------------------------------------------

/// Manages background shell processes, identified by unique IDs.
pub struct ShellManager {
    shells: Mutex<HashMap<String, Arc<Mutex<ManagedShell>>>>,
}

impl Default for ShellManager {
    fn default() -> Self {
        Self::new()
    }
}

impl ShellManager {
    pub fn new() -> Self {
        Self {
            shells: Mutex::new(HashMap::new()),
        }
    }

    /// Spawn a new shell process.
    pub async fn start(&self, cmd: &str, cwd: Option<&str>) -> anyhow::Result<String> {
        let shell_id = format!("bash-{}", uuid::Uuid::new_v4().to_string()[..8].to_owned());
        let mut command = Command::new("sh");
        command.arg("-c").arg(cmd);
        if let Some(dir) = cwd {
            command.current_dir(dir);
        }
        command
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped());

        let mut child = command.spawn()?;

        let stdout = child.stdout.take();
        let stderr = child.stderr.take();

        let shell = ManagedShell {
            shell_id: shell_id.clone(),
            cmd: cmd.to_owned(),
            cwd: cwd.map(|s| s.to_owned()),
            output_chunks: Vec::new(),
            child: Some(child),
            exit_code: None,
            read_handles: Vec::new(),
        };

        let shell = Arc::new(Mutex::new(shell));

        // Spawn tasks to drain stdout and stderr.
        let stdout_shell = shell.clone();
        let stdout_handle = tokio::spawn(async move { drain_stream(stdout_shell, stdout).await });

        let stderr_shell = shell.clone();
        let stderr_handle = tokio::spawn(async move { drain_stream(stderr_shell, stderr).await });

        {
            let mut s = shell.lock().await;
            s.read_handles.push(stdout_handle);
            s.read_handles.push(stderr_handle);
        }

        self.shells.lock().await.insert(shell_id.clone(), shell);

        Ok(shell_id)
    }

    /// Get a snapshot of a shell's current state.
    /// If the shell has finished (exit code present), it is removed from the
    /// map to prevent memory leaks.
    pub async fn get_output(
        &self,
        shell_id: &str,
    ) -> anyhow::Result<(String, Option<i32>, String)> {
        let shell_arc = {
            let shells = self.shells.lock().await;
            shells
                .get(shell_id)
                .ok_or_else(|| anyhow::anyhow!("unknown shell id: {shell_id}"))?
                .clone()
        };
        let shell = shell_arc.lock().await;
        let result = (
            shell.output(),
            shell.returncode(),
            shell.status().to_owned(),
        );
        // Auto-cleanup: if the shell has exited, remove it from the map.
        if result.1.is_some() {
            drop(shell);
            self.shells.lock().await.remove(shell_id);
        }
        Ok(result)
    }

    /// Get the shell ID, output, return code.
    pub async fn get(&self, shell_id: &str) -> anyhow::Result<(String, Option<i32>, String)> {
        self.get_output(shell_id).await
    }

    /// Terminate a shell process.
    pub async fn terminate(&self, shell_id: &str) -> anyhow::Result<(String, Option<i32>, String)> {
        let shells = self.shells.lock().await;
        let shell_arc = shells
            .get(shell_id)
            .ok_or_else(|| anyhow::anyhow!("unknown shell id: {shell_id}"))?
            .clone();
        drop(shells);

        let mut shell = shell_arc.lock().await;
        if shell.exit_code.is_none()
            && let Some(ref mut child) = shell.child
        {
            let _ = child.kill().await;
            let status = child.wait().await.ok();
            shell.exit_code = status.and_then(|s| s.code()).or(Some(-1));
        }
        finalize_shell(&mut shell).await;
        Ok((
            shell.output(),
            shell.returncode(),
            shell.status().to_owned(),
        ))
    }

    /// Wait for a shell to finish.
    pub async fn wait_closed(
        &self,
        shell_id: &str,
    ) -> anyhow::Result<(String, Option<i32>, String)> {
        let shells = self.shells.lock().await;
        let shell_arc = shells
            .get(shell_id)
            .ok_or_else(|| anyhow::anyhow!("unknown shell id: {shell_id}"))?
            .clone();
        drop(shells);

        let mut shell = shell_arc.lock().await;
        if shell.exit_code.is_none()
            && let Some(ref mut child) = shell.child
        {
            let status = child.wait().await.ok();
            shell.exit_code = status.and_then(|s| s.code()).or(Some(-1));
        }
        finalize_shell(&mut shell).await;
        Ok((
            shell.output(),
            shell.returncode(),
            shell.status().to_owned(),
        ))
    }
}

async fn finalize_shell(shell: &mut ManagedShell) {
    let handles: Vec<_> = shell.read_handles.drain(..).collect();
    for handle in handles {
        if let Ok(chunks) = handle.await {
            shell.output_chunks.extend(chunks);
        }
    }
}

async fn drain_stream(
    _shell: Arc<Mutex<ManagedShell>>,
    stream: Option<impl tokio::io::AsyncRead + Unpin>,
) -> Vec<String> {
    let mut chunks = Vec::new();
    let Some(mut reader) = stream else {
        return chunks;
    };
    let mut buf = [0u8; 4096];
    loop {
        match reader.read(&mut buf).await {
            Ok(0) => break,
            Ok(n) => {
                let text = String::from_utf8_lossy(&buf[..n]).to_string();
                chunks.push(text);
            }
            Err(_) => break,
        }
    }
    chunks
}

/// Global shell manager singleton.
static SHELL_MANAGER: std::sync::LazyLock<ShellManager> =
    std::sync::LazyLock::new(ShellManager::new);

/// Access the global shell manager.
pub fn shell_manager() -> &'static ShellManager {
    &SHELL_MANAGER
}

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

    #[tokio::test]
    async fn test_shell_manager_start_and_wait() {
        let mgr = ShellManager::new();
        let shell_id = mgr.start("echo hello", None).await.unwrap();
        assert!(shell_id.starts_with("bash-"));

        let (output, code, status) = mgr.wait_closed(&shell_id).await.unwrap();
        assert!(output.contains("hello"));
        assert_eq!(code, Some(0));
        assert_eq!(status, "exited");
    }

    #[tokio::test]
    async fn test_shell_manager_terminate() {
        let mgr = ShellManager::new();
        let shell_id = mgr.start("sleep 60", None).await.unwrap();

        let (_, code, status) = mgr.terminate(&shell_id).await.unwrap();
        assert!(code.is_some());
        assert_eq!(status, "exited");
    }

    #[tokio::test]
    async fn test_shell_manager_get_unknown_id_returns_error() {
        let mgr = ShellManager::new();
        let result = mgr.get("nonexistent-id").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_shell_manager_start_with_cwd() {
        let mgr = ShellManager::new();
        let shell_id = mgr.start("pwd", Some("/tmp")).await.unwrap();
        let (output, code, _) = mgr.wait_closed(&shell_id).await.unwrap();
        // On macOS /tmp may be a symlink to /private/tmp
        assert!(output.contains("tmp"));
        assert_eq!(code, Some(0));
    }

    #[tokio::test]
    async fn test_shell_manager_get_output() {
        let mgr = ShellManager::new();
        let shell_id = mgr.start("echo test_output", None).await.unwrap();
        let (output, code, status) = mgr.wait_closed(&shell_id).await.unwrap();
        assert!(output.contains("test_output"));
        assert_eq!(code, Some(0));
        assert_eq!(status, "exited");
    }

    #[test]
    fn test_managed_shell_status() {
        let shell = ManagedShell {
            shell_id: "test".into(),
            cmd: "echo hi".into(),
            cwd: None,
            output_chunks: vec!["hello ".into(), "world".into()],
            child: None,
            exit_code: Some(0),
            read_handles: Vec::new(),
        };
        assert_eq!(shell.output(), "hello world");
        assert_eq!(shell.returncode(), Some(0));
        assert_eq!(shell.status(), "exited");
    }

    #[test]
    fn test_managed_shell_running_status() {
        let shell = ManagedShell {
            shell_id: "test".into(),
            cmd: "sleep 10".into(),
            cwd: None,
            output_chunks: Vec::new(),
            child: None,
            exit_code: None,
            read_handles: Vec::new(),
        };
        assert_eq!(shell.status(), "running");
        assert_eq!(shell.returncode(), None);
    }

    #[test]
    fn test_global_shell_manager_exists() {
        // Just verify the global singleton is accessible
        let _mgr = shell_manager();
    }
}