iron-shell 14.0.2

The world's first multi-agent shell with typed functional pipelines and multi-modal AI
use std::collections::{BTreeMap, HashSet};

use rustc_hash::FxHashMap;
use std::path::{Path, PathBuf};

use crate::value::Value;

/// Process/runtime environment for the IronShell evaluator & builtins.
///
/// - Stores variables (including the pipeline slot `__pipe_input__`)
/// - Tracks mutability: variables are immutable by default
/// - Tracks visibility: variables are private by default, `pub` makes them exportable
/// - Optionally tracks a current working directory (cwd)
#[derive(Debug, Default, Clone)]
pub struct Env {
    /// A hash map, not a B-tree: every lambda call binds and unbinds its
    /// parameters here, among ~120 module entries, and the B-tree's
    /// insert/remove (comparisons plus rebalancing) was two thirds of a lambda
    /// `map`. Nothing depends on the order; `vars()` callers sort.
    vars: FxHashMap<String, Value>,
    /// Track which variables are mutable (created with `let mut`)
    mutable_vars: FxHashMap<String, bool>,
    /// Names the *user* declared with `let mut`.
    ///
    /// Distinct from `mutable_vars`, which `set_var_unchecked` also sets for
    /// internal bindings such as lambda parameters and catch variables. Those
    /// are not user-mutable, and conflating the two made a lambda parameter
    /// look reassignable — which mattered once closures needed to know whether
    /// capturing a name by value could go stale.
    declared_mutable: std::collections::BTreeSet<String>,

    /// Track which variables are public (can be imported by other modules)
    public_vars: HashSet<String>,
    /// Explicitly exported names (from `export { ... }` statements)
    exported_names: HashSet<String>,
    cwd: Option<PathBuf>,
    pipe_input: Option<Value>,
}

impl Env {
    /// Create a fresh environment.
    pub fn new() -> Self {
        Self {
            vars: FxHashMap::default(),
            mutable_vars: FxHashMap::default(),
            declared_mutable: std::collections::BTreeSet::new(),
            public_vars: HashSet::new(),
            exported_names: HashSet::new(),
            cwd: None,
            pipe_input: None,
        }
    }

    /// Register a module as an immutable global variable.
    /// Used by the module system to provide namespaced access to builtins.
    pub fn register_module(&mut self, name: &str, module: Value) {
        self.vars.insert(name.to_string(), module);
        self.mutable_vars.insert(name.to_string(), false);
    }

    pub fn input(&self) -> Option<&Value> {
        self.pipe_input.as_ref()
    }
    pub fn set_input(&mut self, v: Option<Value>) {
        self.pipe_input = v;
    }
    pub fn take_input(&mut self) -> Option<Value> {
        self.pipe_input.take()
    }
    /// Clears the current pipeline input value
    pub fn clear_input(&mut self) {
        self.pipe_input = None;
    }

    // -------------------------
    // Variable accessors
    // -------------------------

    /// Immutable lookup. Works on `&Env` (what your error needed).
    pub fn get_var(&self, name: &str) -> Option<&Value> {
        self.vars.get(name)
    }

    /// Is this name one of the pre-bound modules (`math`, `git`, `user`, …)?
    ///
    /// Worth distinguishing, because the generic message sends people looking
    /// for an earlier `let` that does not exist: `let user = 1` fails at the
    /// first line of a fresh script, and "Cannot reassign immutable variable"
    /// describes neither the cause nor the fix.
    fn is_module_binding(&self, name: &str) -> bool {
        // `crate::modules` only exists under the `native` feature; on wasm no
        // modules are pre-bound, so the question cannot arise and the generic
        // message is the right one.
        #[cfg(feature = "native")]
        {
            self.vars.contains_key(name)
                && !self.is_mutable(name)
                && crate::modules::is_module_name(name)
        }
        #[cfg(not(feature = "native"))]
        {
            let _ = name;
            false
        }
    }

    /// The refusal message for binding over an existing immutable name.
    fn immutable_binding_error(&self, name: &str) -> String {
        if self.is_module_binding(name) {
            format!(
                "'{}' is the name of a built-in module and cannot be used as a variable name; pick another name (the module stays reachable as `{}.…`)",
                name, name
            )
        } else {
            format!(
                "Cannot reassign immutable variable '{}'. Use 'let mut {}' to make it mutable.",
                name, name
            )
        }
    }

    /// The hint for a `let` that cannot bind `name` because the name is taken.
    /// It is the statement's own advice, so it fits the refusal: a module name
    /// is not something to `let mut` over.
    pub fn rebind_hint(&self, name: &str) -> String {
        if self.is_module_binding(name) {
            format!(
                "`{name}` is a module name; use a different name (the module stays reachable as `{name}.…`)"
            )
        } else {
            format!("bind `{name}` once with `let mut {name} = …`, or use a different name")
        }
    }

    /// The hint for a bare `name = value` that cannot reassign `name`: the
    /// binding has to be mutable from its first `let`.
    pub fn reassign_hint(&self, name: &str) -> String {
        if self.is_module_binding(name) {
            self.rebind_hint(name)
        } else {
            format!(
                "bind `{name}` mutable from the start, as `mut {name} = 0` or `let mut {name} = 0`, or use a different name"
            )
        }
    }

    /// Set a new variable or reassign if mutable.
    /// Returns error if trying to reassign an immutable variable.
    pub fn set_var<S: Into<String>>(&mut self, name: S, value: Value) -> Result<(), String> {
        let name_str = name.into();

        // Check if variable already exists
        if self.vars.contains_key(&name_str) {
            // Check mutability
            if !self.is_mutable(&name_str) {
                return Err(self.immutable_binding_error(&name_str));
            }
        }

        self.vars.insert(name_str, value);
        Ok(())
    }

    /// Declare a new variable with specified mutability.
    /// This is used for initial `let` or `let mut` declarations.
    pub fn declare_var<S: Into<String>>(
        &mut self,
        name: S,
        value: Value,
        is_mut: bool,
    ) -> Result<(), String> {
        let name_str = name.into();

        // Check if variable already exists
        if self.vars.contains_key(&name_str) {
            // If existing variable is immutable, prevent shadowing/reassignment
            if !self.is_mutable(&name_str) {
                return Err(self.immutable_binding_error(&name_str));
            }
            // If mutable, allow update but preserve mutability
            self.vars.insert(name_str, value);
            return Ok(());
        }

        // New variable - set value and mutability
        self.vars.insert(name_str.clone(), value);
        if is_mut {
            self.declared_mutable.insert(name_str.clone());
        } else {
            self.declared_mutable.remove(&name_str);
        }
        self.mutable_vars.insert(name_str, is_mut);
        Ok(())
    }

    /// Internal use only: Set a variable without mutability checks.
    /// Used for pattern matching, lambda params, builtins, etc.
    /// Bind a lambda parameter. Like `set_var_unchecked`, but allocates the
    /// name once instead of twice, and marks it mutable only the first time:
    /// a lambda call binds and unbinds its parameters per element, so this is
    /// on the hottest path the evaluator has.
    #[cfg_attr(not(feature = "native"), allow(dead_code))] // only the evaluator calls it
    pub(crate) fn bind_param(&mut self, name: &str, value: Value) {
        self.vars.insert(name.to_string(), value);
        if !self.mutable_vars.contains_key(name) {
            self.mutable_vars.insert(name.to_string(), true);
        }
    }

    /// Replace a live parameter without reallocating its key. Collection
    /// builtins keep these slots alive for the duration of their loop.
    #[cfg_attr(not(feature = "native"), allow(dead_code))] // only the evaluator calls it
    pub(crate) fn replace_param(&mut self, name: &str, value: Value) -> Option<Value> {
        if let Some(slot) = self.vars.get_mut(name) {
            Some(std::mem::replace(slot, value))
        } else {
            self.bind_param(name, value);
            None
        }
    }

    #[cfg_attr(not(feature = "native"), allow(dead_code))] // only the evaluator calls it
    pub(crate) fn restore_param(&mut self, name: &str, previous: Option<Value>) -> Option<Value> {
        match previous {
            Some(value) => self.replace_param(name, value),
            None => self.take_var(name),
        }
    }

    #[cfg_attr(not(feature = "native"), allow(dead_code))] // only the evaluator calls it
    pub(crate) fn set_var_unchecked<S: Into<String>>(&mut self, name: S, value: Value) {
        let name_str = name.into();
        self.vars.insert(name_str.clone(), value);
        // Mark as mutable for internal bindings
        self.mutable_vars.insert(name_str, true);
    }

    /// Did the user declare this name with `let mut`?
    ///
    /// `is_mutable` is true for internal bindings too, so it cannot answer this.
    pub fn is_declared_mutable(&self, name: &str) -> bool {
        self.declared_mutable.contains(name)
    }

    /// Check if a variable is mutable
    pub fn is_mutable(&self, name: &str) -> bool {
        self.mutable_vars.get(name).copied().unwrap_or(false)
    }

    /// Check if a variable is public (can be imported)
    pub fn is_public(&self, name: &str) -> bool {
        self.public_vars.contains(name)
    }

    /// Check if a variable is exported (either pub or explicitly exported)
    pub fn is_exported(&self, name: &str) -> bool {
        self.public_vars.contains(name) || self.exported_names.contains(name)
    }

    /// Mark a variable as public
    pub fn set_public(&mut self, name: &str) {
        self.public_vars.insert(name.to_string());
    }

    /// Add an explicit export
    pub fn add_export(&mut self, name: &str) {
        self.exported_names.insert(name.to_string());
    }

    /// Get all exported variable names (pub vars + explicit exports)
    pub fn exports(&self) -> impl Iterator<Item = &String> {
        self.public_vars.iter().chain(self.exported_names.iter())
    }

    /// Get exported variables as a map (for import "path" as mod)
    pub fn exported_vars(&self) -> BTreeMap<String, Value> {
        let mut exports = BTreeMap::new();
        for name in self.exports() {
            if let Some(value) = self.vars.get(name) {
                exports.insert(name.clone(), value.clone());
            }
        }
        exports
    }

    /// Delete a variable if present.
    pub fn del_var(&mut self, name: &str) {
        self.vars.remove(name);
    }

    /// Mutable lookup (if you ever need to mutate a stored value in place).
    pub fn get_var_mut(&mut self, name: &str) -> Option<&mut Value> {
        self.vars.get_mut(name)
    }

    /// Take (remove and return) a variable. Handy for temporary slots.
    pub fn take_var(&mut self, name: &str) -> Option<Value> {
        self.vars.remove(name)
    }

    /// Expose the whole map if some builtin needs to iterate (read-only).
    pub fn vars(&self) -> &FxHashMap<String, Value> {
        &self.vars
    }

    /// Expose mutable map (use sparingly).
    pub fn vars_mut(&mut self) -> &mut FxHashMap<String, Value> {
        &mut self.vars
    }

    // -------------------------
    // Working directory helpers
    // -------------------------

    /// Get current working directory tracked by the shell (if any).
    pub fn cwd(&self) -> Option<&Path> {
        self.cwd.as_deref()
    }

    /// Set current working directory tracked by the shell.
    pub fn set_cwd<P: Into<PathBuf>>(&mut self, p: P) {
        self.cwd = Some(p.into());
    }
}