//! Python script executor with shell-escape path removal.
//!
//! Spawns a Python process (`python3` or `micropython`) and injects a safety
//! preamble that removes dangerous `os` attributes and blocks shell-related
//! modules. The preamble also installs a custom `__import__` guard to prevent
//! bypasses via `importlib` or `__builtins__.__import__`.
//!
//! # Defense layers
//!
//! When the engine is `python3` (detected by binary name starting with
//! `"python"`), a preamble is prepended to every script. The preamble applies
//! three independent defense layers:
//!
//! ## Layer 1 — `sys.modules` poisoning
//!
//! ```python
//! sys.modules['subprocess'] = None
//! sys.modules['shutil'] = None
//! # ... (all DEFAULT_DENY + user-specified PYEXEC_DENY)
//! ```
//!
//! Setting a module to `None` in `sys.modules` causes `import subprocess` to
//! raise `ImportError`. This is the primary block and catches all standard
//! `import` statements.
//!
//! ## Layer 2 — `__import__` guard
//!
//! ```python
//! def _safe_import(name, *args, _deny=_deny_set, _imp=_orig_import, **kwargs):
//!     if name in _deny or name.split('.')[0] in _deny:
//!         raise ImportError(...)
//!     return _imp(name, *args, **kwargs)
//! __builtins__.__import__ = _safe_import
//! ```
//!
//! Replaces the built-in `__import__` function with a wrapper that checks
//! against the deny set. This prevents bypasses via:
//!
//! - `__builtins__.__import__('subprocess')` — direct call
//! - `importlib.import_module('subprocess')` — `importlib` itself is in the
//!   deny list (Layer 1), but even if somehow loaded, the underlying
//!   `__import__` is guarded
//! - `importlib.reload(os)` — same: `importlib` is blocked at Layer 1
//!
//! Default arguments (`_deny=_deny_set, _imp=_orig_import`) capture values
//! at definition time, so subsequent `del` of the preamble variables does not
//! break the guard.
//!
//! ## Layer 3 — `os` attribute removal
//!
//! ```python
//! for attr in ('system', 'popen', 'exec', 'execv', 'spawn*', 'fork', 'kill', ...):
//!     delattr(os, attr)
//! ```
//!
//! Removes shell-execution functions from the `os` module object in-place.
//! Since `import os` returns the cached module object, subsequent imports in
//! user code see the stripped version.
//!
//! ## Optional — `open()` path guard (deny_paths)
//!
//! When [`ExecConfig::deny_paths`] is non-empty, an additional guard replaces
//! the built-in `open()`:
//!
//! ```python
//! def _safe_open(file, *args, _blocked=_deny_paths, ...):
//!     resolved = os.path.realpath(os.path.abspath(str(file)))
//!     for bp in _blocked:
//!         if resolved == bp or resolved.startswith(bp + '/'):
//!             raise PermissionError(...)
//!     return _orig_open(file, *args, **kwargs)
//! ```
//!
//! Both the blocked paths and the target path are resolved via `realpath()`
//! to handle symlinks (e.g. macOS `/tmp` → `/private/tmp`).
//!
//! # Known limitations
//!
//! - **Not a sandbox.** File I/O (except `deny_paths`), network (`socket`,
//!   `http.client`), and memory are unrestricted. For full isolation, use a
//!   container or VM.
//! - **Transitive dependencies.** Some stdlib modules import denied modules
//!   internally. For example, `tempfile` imports `shutil`, so
//!   `import tempfile` will fail when `shutil` is in the deny list (it is by
//!   default). This is intentional — `shutil` contains `rmtree` and other
//!   destructive filesystem operations.
//! - **C extensions.** A pre-loaded C extension could call `exec()` at the C
//!   level, bypassing Python-layer guards. This is out of scope for a
//!   preamble-based approach.
//! - **`micropython`** does not receive the preamble (these APIs do not exist
//!   in micropython).

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

use tokio::process::Command;

/// Errors returned by [`execute`] and [`execute_file`].
#[derive(Debug, thiserror::Error)]
pub enum ExecError {
    #[error("python binary not found: '{0}'. Install python3 or micropython")]
    BinaryNotFound(String),
    #[error("execution timed out after {0} seconds")]
    Timeout(u64),
    #[error("execution failed: {0}")]
    ProcessError(String),
    #[error("io error: {0}")]
    Io(#[from] std::io::Error),
}

// =============================================================================
// ExecConfig
// =============================================================================

/// Configuration for Python script execution.
///
/// Constructed from environment variables via [`ExecConfig::from_env`], or
/// built manually for testing. Fields can be overridden at runtime through
/// the MCP `session_start` tool.
#[derive(Debug, Clone)]
pub struct ExecConfig {
    /// Python binary (python3, micropython, or absolute path).
    pub binary: String,
    /// Maximum execution time.
    pub timeout: Duration,
    /// Working directory for script execution.
    pub workdir: Option<PathBuf>,
    /// Additional modules to block (beyond the hardcoded defaults).
    pub deny_modules: Vec<String>,
    /// When true, `execute_file` restricts paths to within `workdir`.
    pub restrict_path: bool,
    /// Paths to block from `open()`. Prefix-matched after canonicalization.
    pub deny_paths: Vec<PathBuf>,
}

/// Default modules always blocked when running under python3.
const DEFAULT_DENY: &[&str] = &[
    "subprocess",
    "shutil",
    "ctypes",
    "multiprocessing",
    "signal",
    "importlib",
    "importlib.util",
    "importlib.abc",
    "importlib.machinery",
    "importlib.resources",
];

impl ExecConfig {
    /// Build config from ENV → detect.
    ///
    /// ENV precedence:
    /// - `PYEXEC_ENGINE`: binary name/path (skips auto-detect)
    /// - `PYEXEC_TIMEOUT`: seconds (default 30, max 300)
    /// - `PYEXEC_WORKDIR`: default working directory
    /// - `PYEXEC_DENY`: comma-separated additional modules to block
    /// - `PYEXEC_RESTRICT_PATH`: "true"/"1" to restrict execute_file to workdir (default: off)
    pub fn from_env() -> Self {
        let binary = std::env::var("PYEXEC_ENGINE")
            .ok()
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| detect_binary().unwrap_or_else(|| "micropython".to_string()));

        let timeout_secs = std::env::var("PYEXEC_TIMEOUT")
            .ok()
            .and_then(|s| s.parse::<u64>().ok())
            .map(|s| s.min(300))
            .unwrap_or(30);

        let workdir = std::env::var("PYEXEC_WORKDIR")
            .ok()
            .filter(|s| !s.is_empty())
            .map(PathBuf::from)
            .filter(|p| p.is_dir());

        let deny_modules = std::env::var("PYEXEC_DENY")
            .ok()
            .filter(|s| !s.is_empty())
            .map(|s| s.split(',').map(|m| m.trim().to_string()).collect())
            .unwrap_or_default();

        let restrict_path = std::env::var("PYEXEC_RESTRICT_PATH")
            .ok()
            .is_some_and(|v| matches!(v.as_str(), "true" | "1"));

        let deny_paths = std::env::var("PYEXEC_DENY_PATHS")
            .ok()
            .filter(|s| !s.is_empty())
            .map(|s| {
                s.split(':')
                    .map(|p| {
                        let expanded = if p.starts_with('~') {
                            if let Ok(home) = std::env::var("HOME") {
                                p.replacen('~', &home, 1)
                            } else {
                                p.to_string()
                            }
                        } else {
                            p.to_string()
                        };
                        PathBuf::from(expanded)
                    })
                    .collect()
            })
            .unwrap_or_default();

        Self {
            binary,
            timeout: Duration::from_secs(timeout_secs),
            workdir,
            deny_modules,
            restrict_path,
            deny_paths,
        }
    }

    /// All denied modules: defaults + env extras.
    fn all_deny_modules(&self) -> Vec<String> {
        let mut modules: Vec<String> = DEFAULT_DENY.iter().map(|s| s.to_string()).collect();
        for m in &self.deny_modules {
            if !modules.contains(m) {
                modules.push(m.clone());
            }
        }
        modules
    }
}

impl Default for ExecConfig {
    fn default() -> Self {
        Self::from_env()
    }
}

// =============================================================================
// Binary detection
// =============================================================================

/// Check which Python binary is available on PATH. Prefers python3.
fn detect_binary() -> Option<String> {
    for candidate in &["python3", "micropython"] {
        if std::process::Command::new(candidate)
            .arg("--version")
            .stdout(std::process::Stdio::null())
            .stderr(std::process::Stdio::null())
            .status()
            .is_ok_and(|s| s.success())
        {
            return Some(candidate.to_string());
        }
    }
    None
}

// =============================================================================
// Sandbox preamble
// =============================================================================

/// OS attributes always removed when running under python3.
const DANGEROUS_OS_ATTRS: &[&str] = &[
    "system",
    "popen",
    "exec",
    "execl",
    "execle",
    "execlp",
    "execlpe",
    "execv",
    "execve",
    "execvp",
    "execvpe",
    "spawnl",
    "spawnle",
    "spawnlp",
    "spawnlpe",
    "spawnv",
    "spawnve",
    "spawnvp",
    "spawnvpe",
    "posix_spawn",
    "posix_spawnp",
    "fork",
    "forkpty",
    "kill",
    "killpg",
    "plock",
    "startfile",
];

fn needs_sandbox(binary: &str) -> bool {
    let name = std::path::Path::new(binary)
        .file_name()
        .and_then(|n| n.to_str())
        .unwrap_or(binary);
    name.starts_with("python")
}

fn build_preamble(config: &ExecConfig) -> String {
    let deny = config.all_deny_modules();
    let module_blocks: String = deny
        .iter()
        .map(|m| format!("_sys.modules['{m}'] = None"))
        .collect::<Vec<_>>()
        .join("\n");

    let os_attrs: String = DANGEROUS_OS_ATTRS
        .iter()
        .map(|a| format!("'{a}'"))
        .collect::<Vec<_>>()
        .join(", ");

    let deny_set: String = deny
        .iter()
        .map(|m| format!("'{m}'"))
        .collect::<Vec<_>>()
        .join(", ");

    let open_guard = if config.deny_paths.is_empty() {
        String::new()
    } else {
        let paths: String = config
            .deny_paths
            .iter()
            .map(|p| format!("'{}'", p.display()))
            .collect::<Vec<_>>()
            .join(", ");
        format!(
            r#"
import os.path as _osp
_orig_open = open
_deny_paths = tuple(_osp.realpath(p) for p in ({paths},))
def _safe_open(file, *args, _blocked=_deny_paths, _op=_orig_open, _rp=_osp.realpath, _absp=_osp.abspath, **kwargs):
    _resolved = _rp(_absp(str(file)))
    for _bp in _blocked:
        if _resolved == _bp or _resolved.startswith(_bp + '/'):
            raise PermissionError("access to '{{0}}' is blocked".format(file))
    return _op(file, *args, **kwargs)
if hasattr(__builtins__, 'open'):
    __builtins__.open = _safe_open
else:
    __builtins__['open'] = _safe_open
del _osp, _orig_open, _deny_paths, _safe_open
"#
        )
    };

    format!(
        r#"import sys as _sys
{module_blocks}
import os as _os
for _attr in ({os_attrs}):
    if hasattr(_os, _attr):
        delattr(_os, _attr)
_orig_import = __builtins__.__import__ if hasattr(__builtins__, '__import__') else __import__
_deny_set = frozenset(({deny_set}))
def _safe_import(name, *args, _deny=_deny_set, _imp=_orig_import, **kwargs):
    if name in _deny or name.split('.')[0] in _deny:
        raise ImportError("module '{{0}}' is blocked".format(name))
    return _imp(name, *args, **kwargs)
if hasattr(__builtins__, '__import__'):
    __builtins__.__import__ = _safe_import
else:
    __builtins__['__import__'] = _safe_import
del _os, _sys, _orig_import, _deny_set, _safe_import{open_guard}
"#
    )
}

fn sandboxed_script(script: &str, config: &ExecConfig) -> String {
    if needs_sandbox(&config.binary) {
        let preamble = build_preamble(config);
        format!("{preamble}\n{script}")
    } else {
        script.to_string()
    }
}

// =============================================================================
// Execution
// =============================================================================

/// Output captured from a Python script execution.
#[derive(Debug, serde::Serialize)]
pub struct ExecResult {
    pub stdout: String,
    pub stderr: String,
    pub exit_code: i32,
    pub timed_out: bool,
    /// `true` if the process was terminated by a signal (Unix only).
    pub signaled: bool,
}

/// Execute a Python script string and return captured output.
///
/// When running under `python3`, a safety preamble is prepended that removes
/// shell-escape functions and blocks denied modules.
pub async fn execute(script: &str, config: &ExecConfig) -> Result<ExecResult, ExecError> {
    let full_script = sandboxed_script(script, config);
    let mut cmd = Command::new(&config.binary);
    cmd.arg("-c").arg(&full_script);
    spawn_and_wait(cmd, config).await
}

/// Execute a Python script file.
///
/// For `python3`, the file contents are read and executed via `-c` with the
/// safety preamble prepended (so the preamble applies even to file execution).
/// When `restrict_path` is enabled, `path` must be within `workdir`.
pub async fn execute_file(path: &Path, config: &ExecConfig) -> Result<ExecResult, ExecError> {
    if !path.exists() {
        return Err(ExecError::ProcessError(format!(
            "file not found: {}",
            path.display()
        )));
    }

    if config.restrict_path {
        validate_path(path, config)?;
    }

    let mut cmd = Command::new(&config.binary);
    if needs_sandbox(&config.binary) {
        let content = std::fs::read_to_string(path).map_err(|e| {
            ExecError::ProcessError(format!("failed to read {}: {e}", path.display()))
        })?;
        let full_script = sandboxed_script(&content, config);
        cmd.arg("-c").arg(&full_script);
    } else {
        cmd.arg(path);
    }
    spawn_and_wait(cmd, config).await
}

/// Spawn a command and wait for output with timeout.
/// Sets `kill_on_drop(true)` so the child process is killed if the future is cancelled.
async fn spawn_and_wait(mut cmd: Command, config: &ExecConfig) -> Result<ExecResult, ExecError> {
    cmd.stdin(Stdio::null());
    cmd.stdout(Stdio::piped());
    cmd.stderr(Stdio::piped());

    if let Some(ref wd) = config.workdir {
        cmd.current_dir(wd);
    }

    cmd.kill_on_drop(true);

    let child = cmd.spawn().map_err(|e| {
        if e.kind() == std::io::ErrorKind::NotFound {
            ExecError::BinaryNotFound(config.binary.clone())
        } else {
            ExecError::Io(e)
        }
    })?;

    match tokio::time::timeout(config.timeout, child.wait_with_output()).await {
        Ok(Ok(output)) => {
            let stdout = String::from_utf8_lossy(&output.stdout).to_string();
            let stderr = String::from_utf8_lossy(&output.stderr).to_string();
            let signaled = output.status.code().is_none();
            let exit_code = output.status.code().unwrap_or(128);
            Ok(ExecResult {
                stdout,
                stderr,
                exit_code,
                timed_out: false,
                signaled,
            })
        }
        Ok(Err(e)) => Err(ExecError::ProcessError(e.to_string())),
        Err(_) => Err(ExecError::Timeout(config.timeout.as_secs())),
    }
}

/// Validate that `path` is within `workdir` (canonicalized).
fn validate_path(path: &Path, config: &ExecConfig) -> Result<(), ExecError> {
    let Some(ref workdir) = config.workdir else {
        return Err(ExecError::ProcessError(
            "restrict_path is enabled but workdir is not set".to_string(),
        ));
    };
    let canon_workdir = workdir.canonicalize().map_err(|e| {
        ExecError::ProcessError(format!("cannot resolve workdir {}: {e}", workdir.display()))
    })?;
    let canon_path = path.canonicalize().map_err(|e| {
        ExecError::ProcessError(format!("cannot resolve path {}: {e}", path.display()))
    })?;
    if !canon_path.starts_with(&canon_workdir) {
        return Err(ExecError::ProcessError(format!(
            "path {} is outside workdir {}",
            canon_path.display(),
            canon_workdir.display()
        )));
    }
    Ok(())
}
