//! # Shell Command Execution
//!
//! Executes shell commands in a working directory (async).
//! Provides timeout protection and output truncation to prevent OOM.
//!
//! ## Safety Design (双层截断)
//!
//! - **Shell layer**: Truncates output at configurable limit (default 1MB) to prevent OOM
//! - **Memory layer**: Further truncation handled by caller if needed

use std::path::{Path, PathBuf};
use std::time::Duration;

use thiserror::Error;

// ---------------------------------------------------------------------------
// Error types
// ---------------------------------------------------------------------------

#[derive(Debug, Error)]
pub enum ShellError {
    #[error("Shell command timed out after {0:?}")]
    Timeout(Duration),

    #[error("Shell I/O error: {0}")]
    Io(String),

    #[error("Shell output not valid UTF-8: {0}")]
    Utf8(String),
}

impl From<std::io::Error> for ShellError {
    fn from(e: std::io::Error) -> Self {
        ShellError::Io(e.to_string())
    }
}

impl From<std::string::FromUtf8Error> for ShellError {
    fn from(e: std::string::FromUtf8Error) -> Self {
        ShellError::Utf8(e.to_string())
    }
}

pub type ShellResult<T> = std::result::Result<T, ShellError>;

// ---------------------------------------------------------------------------
// Output
// ---------------------------------------------------------------------------

/// Result of a shell command execution.
#[derive(Debug, Clone)]
pub struct ShellOutput {
    /// Combined stdout + stderr output (truncated if over limit).
    pub output: String,
    /// Process exit code (None if killed/timeout).
    pub exit_code: Option<i32>,
    /// Whether the output was truncated.
    pub truncated: bool,
    /// Execution time.
    pub duration: Duration,
}

impl ShellOutput {
    /// Whether the command exited successfully (exit code 0).
    pub fn success(&self) -> bool {
        self.exit_code == Some(0)
    }

    /// Human-readable exit code for error messages.
    pub fn exit_code_display(&self) -> String {
        self.exit_code
            .map_or("unknown".to_string(), |c| c.to_string())
    }
}

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

/// Default maximum output size in bytes (1MB).
const DEFAULT_MAX_OUTPUT: usize = 1024 * 1024;

/// Default command timeout in seconds.
const DEFAULT_TIMEOUT_SECS: u64 = 300;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Async read from `reader` into `buf` until EOF or `limit` bytes accumulated.
/// Uses chunked reads to avoid closing the pipe prematurely (SIGPIPE fix).
async fn read_limited_async(
    reader: &mut (impl tokio::io::AsyncRead + Unpin),
    limit: usize,
    buf: &mut Vec<u8>,
) {
    use tokio::io::AsyncReadExt;
    let mut chunk = [0u8; 8192];
    while buf.len() < limit {
        let remaining = limit - buf.len();
        let to_read = remaining.min(chunk.len());
        match reader.read(&mut chunk[..to_read]).await {
            Ok(0) => break,
            Ok(n) => buf.extend_from_slice(&chunk[..n]),
            Err(_) => break,
        }
    }
}

// ---------------------------------------------------------------------------
// Shell
// ---------------------------------------------------------------------------

/// Shell command executor with timeout and truncation protection.
///
/// Both privileged and sandbox modes pipe scripts via stdin to bash,
/// ensuring identical I/O behavior.
///
/// - No sandbox_user (privileged): `/bin/bash` with `current_dir`
/// - With sandbox_user: `su - {user} -c "bash -s"` (紧箍咒)
pub(crate) struct Shell {
    working_dir: PathBuf,
    timeout_duration: Duration,
    max_output: usize,
    /// Linux username for sandboxed execution (紧箍咒)
    sandbox_user: Option<String>,
}

impl Shell {
    pub fn new(working_dir: PathBuf, sandbox_user: Option<String>) -> Self {
        Self {
            working_dir,
            timeout_duration: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
            max_output: DEFAULT_MAX_OUTPUT,
            sandbox_user,
        }
    }

    /// Detect whether a sandbox user exists for the given instance ID.
    /// Returns `Some("agent-{id}")` if the system user exists, `None` otherwise.
    pub fn detect_sandbox_user(instance_id: &str) -> Option<String> {
        let user = format!("agent-{}", instance_id);
        let exists = std::process::Command::new("id")
            .arg(&user)
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::null())
            .status()
            .map(|s| s.success())
            .unwrap_or(false);
        if exists {
            Some(user)
        } else {
            None
        }
    }

    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout_duration = timeout;
        self
    }

    pub fn with_max_output(mut self, max: usize) -> Self {
        self.max_output = max;
        self
    }

    /// Execute a shell script in the configured working directory.
    pub async fn exec(&self, script: &str) -> ShellResult<ShellOutput> {
        self.exec_in_dir(script, &self.working_dir).await
    }

    /// Execute a shell script in a specific directory (async).
    ///
    /// Both paths pipe the script via stdin to bash, ensuring identical I/O behavior.
    /// Sandboxed mode wraps with `su -` for user isolation (紧箍咒).
    pub async fn exec_in_dir(&self, script: &str, dir: &Path) -> ShellResult<ShellOutput> {
        use std::process::Stdio;
        use tokio::process::Command;

        let start = std::time::Instant::now();
        let max_output = self.max_output;
        let timeout = self.timeout_duration;

        let mut child = if self.sandbox_user.is_some() {
            let user = self.sandbox_user.as_deref().unwrap_or("nobody");
            let mut cmd = Command::new("su");
            cmd.arg("-")
                .arg(user)
                .arg("-c")
                .arg("bash -s")
                .env("ALICE_ENGINE_CHILD", "1")
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped());
            cmd.spawn()?
        } else {
            let mut cmd = Command::new("/bin/bash");
            cmd.current_dir(dir)
                .env("ALICE_ENGINE_CHILD", "1")
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .stderr(Stdio::piped());
            cmd.spawn()?
        };

        // Write script to stdin
        {
            let script_content = if self.sandbox_user.is_some() {
                format!("cd {} && {}", dir.display(), script)
            } else {
                script.to_string()
            };
            if let Some(mut stdin) = child.stdin.take() {
                use tokio::io::AsyncWriteExt;
                let _ = stdin.write_all(script_content.as_bytes()).await;
            }
        }

        // Async read with timeout
        let read_and_wait = async {
            let mut combined = Vec::new();

            if let Some(mut stdout) = child.stdout.take() {
                read_limited_async(&mut stdout, max_output, &mut combined).await;
            }

            if let Some(mut stderr) = child.stderr.take() {
                let remaining = max_output.saturating_sub(combined.len());
                if remaining > 0 {
                    read_limited_async(&mut stderr, remaining, &mut combined).await;
                }
            }

            let status = child.wait().await;
            (combined, status)
        };

        match tokio::time::timeout(timeout, read_and_wait).await {
            Ok((raw_output, status)) => {
                let truncated = raw_output.len() >= max_output;
                let output = String::from_utf8_lossy(&raw_output).to_string();
                let exit_code = status.ok().and_then(|s| s.code());
                let duration = start.elapsed();

                Ok(ShellOutput {
                    output,
                    exit_code,
                    truncated,
                    duration,
                })
            }
            Err(_) => {
                let _ = child.kill().await;
                Err(ShellError::Timeout(timeout))
            }
        }
    }

    /// Read a file via shell `cat` command.
    /// Useful for sandboxed instances that lack direct filesystem access.
    pub async fn read_file(&self, path: &str) -> ShellResult<ShellOutput> {
        let escaped = path.replace('\'', "'\\''");
        self.exec(&format!("cat '{}'", escaped)).await
    }

    /// Write content to a file via shell heredoc.
    /// Creates parent directories if needed.
    pub async fn write_file(&self, path: &str, content: &str) -> ShellResult<ShellOutput> {
        let escaped = path.replace('\'', "'\\''");
        let delim = format!(
            "HEREDOC_{}",
            uuid::Uuid::new_v4().to_string().replace('-', "")
        );
        self.exec(&format!(
            "mkdir -p \"$(dirname '{}')\" && cat > '{}' << '{}'\n{}\n{}",
            escaped, escaped, delim, content, delim,
        ))
        .await
    }
}

/// Auto-detect shell environment from system info.
///
/// Returns a human-readable string like "Linux系统（x86_64），请生成bash脚本".
pub fn detect_shell_env() -> String {
    let os = std::process::Command::new("uname")
        .arg("-s")
        .output()
        .ok()
        .and_then(|o| String::from_utf8(o.stdout).ok())
        .map(|s| s.trim().to_string())
        .unwrap_or_else(|| "Linux".to_string());

    let arch = std::process::Command::new("uname")
        .arg("-m")
        .output()
        .ok()
        .and_then(|o| String::from_utf8(o.stdout).ok())
        .map(|s| s.trim().to_string())
        .unwrap_or_else(|| "x86_64".to_string());

    format!("{}系统（{}），请生成bash脚本", os, arch)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn setup() -> (Shell, TempDir) {
        let tmp = TempDir::new().unwrap();
        let shell = Shell::new(tmp.path().to_path_buf(), None);
        (shell, tmp)
    }

    #[tokio::test]
    async fn test_echo() {
        let (shell, _tmp) = setup();
        let result = shell.exec("echo hello").await.unwrap();
        assert_eq!(result.output.trim(), "hello");
        assert_eq!(result.exit_code, Some(0));
        assert!(!result.truncated);
    }

    #[tokio::test]
    async fn test_exit_code() {
        let (shell, _tmp) = setup();
        let result = shell.exec("exit 42").await.unwrap();
        assert_eq!(result.exit_code, Some(42));
    }

    #[tokio::test]
    async fn test_stderr_captured() {
        let (shell, _tmp) = setup();
        let result = shell.exec("echo err >&2").await.unwrap();
        assert!(result.output.contains("err"));
    }

    #[tokio::test]
    async fn test_working_dir() {
        let (shell, tmp) = setup();
        std::fs::write(tmp.path().join("test.txt"), "content").unwrap();
        let result = shell.exec("cat test.txt").await.unwrap();
        assert_eq!(result.output.trim(), "content");
    }

    #[tokio::test]
    async fn test_timeout() {
        let tmp = TempDir::new().unwrap();
        let shell =
            Shell::new(tmp.path().to_path_buf(), None).with_timeout(Duration::from_millis(500));
        let result = shell.exec("sleep 10").await;
        assert!(matches!(result, Err(ShellError::Timeout(_))));
    }

    #[tokio::test]
    async fn test_multiline_script() {
        let (shell, _tmp) = setup();
        let result = shell.exec("echo line1\necho line2").await.unwrap();
        assert!(result.output.contains("line1"));
        assert!(result.output.contains("line2"));
    }

    #[tokio::test]
    async fn test_truncation() {
        let tmp = TempDir::new().unwrap();
        let shell = Shell::new(tmp.path().to_path_buf(), None).with_max_output(100);
        let result = shell.exec("yes | head -200").await.unwrap();
        assert!(result.truncated);
        assert!(result.output.len() <= 200);
    }

    #[tokio::test]
    async fn test_success_helper() {
        let (shell, _tmp) = setup();
        let ok = shell.exec("true").await.unwrap();
        assert!(ok.success());
        let fail = shell.exec("false").await.unwrap();
        assert!(!fail.success());
    }

    #[tokio::test]
    async fn test_exit_code_display() {
        let (shell, _tmp) = setup();
        let result = shell.exec("exit 42").await.unwrap();
        assert_eq!(result.exit_code_display(), "42");
    }

    #[tokio::test]
    async fn test_exit_code_display_unknown() {
        let output = ShellOutput {
            output: String::new(),
            exit_code: None,
            truncated: false,
            duration: Duration::from_secs(0),
        };
        assert_eq!(output.exit_code_display(), "unknown");
    }
}