use crate::registry::EvalRegistry;
use crate::value::Value;
use hamelin_lib::tree::ast::identifier::SimpleIdentifier;
use std::collections::HashMap;
use std::sync::Arc;

/// Environment for variable bindings during evaluation
#[derive(Debug, Clone)]
pub struct Environment {
    /// Variable bindings
    bindings: HashMap<SimpleIdentifier, Value>,
    /// Registry for function evaluation
    registry: Arc<EvalRegistry>,
}

impl Environment {
    /// Create a new empty environment with the default eval registry
    pub fn new() -> Self {
        Self {
            bindings: HashMap::new(),
            registry: Arc::new(EvalRegistry::default()),
        }
    }

    /// Create a new environment with a custom registry
    pub fn with_registry(registry: Arc<EvalRegistry>) -> Self {
        Self {
            bindings: HashMap::new(),
            registry,
        }
    }

    /// Get the eval registry
    pub fn registry(&self) -> &EvalRegistry {
        &self.registry
    }

    /// Bind a variable to a value in this environment
    pub fn bind(&mut self, name: SimpleIdentifier, value: Value) {
        self.bindings.insert(name, value);
    }

    /// Look up a variable in this environment
    pub fn lookup(&self, name: &SimpleIdentifier) -> Option<&Value> {
        self.bindings.get(name)
    }

    /// Check if a variable exists in this environment
    pub fn contains(&self, name: &SimpleIdentifier) -> bool {
        self.bindings.contains_key(name)
    }

    /// Get all bindings
    pub fn bindings(&self) -> &HashMap<SimpleIdentifier, Value> {
        &self.bindings
    }

    /// Clear all bindings
    pub fn clear(&mut self) {
        self.bindings.clear();
    }

    /// Remove a specific binding
    pub fn unbind(&mut self, name: &SimpleIdentifier) -> Option<Value> {
        self.bindings.remove(name)
    }

    /// Get the number of bindings
    pub fn len(&self) -> usize {
        self.bindings.len()
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.bindings.is_empty()
    }

    /// Merge another environment into this one (other environment shadows this one)
    pub fn merge(&mut self, other: Environment) {
        self.bindings.extend(other.bindings);
    }

    /// Create a new environment with the union of two environments
    /// In case of conflicts, the second environment's bindings take precedence
    pub fn union(mut self, other: Environment) -> Self {
        self.merge(other);
        self
    }
}

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

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

    #[test]
    fn test_basic_binding() {
        let mut env = Environment::new();
        let x = SimpleIdentifier::new("x");
        let value = Value::Int(42);

        env.bind(x.clone(), value.clone());
        assert_eq!(env.lookup(&x), Some(&value));
    }

    #[test]
    fn test_merge() {
        let mut env1 = Environment::new();
        env1.bind(SimpleIdentifier::new("a"), Value::Int(1));
        env1.bind(SimpleIdentifier::new("b"), Value::Int(2));

        let mut env2 = Environment::new();
        env2.bind(SimpleIdentifier::new("b"), Value::Int(20));
        env2.bind(SimpleIdentifier::new("c"), Value::Int(3));

        env1.merge(env2);

        assert_eq!(
            env1.lookup(&SimpleIdentifier::new("a")),
            Some(&Value::Int(1))
        );
        assert_eq!(
            env1.lookup(&SimpleIdentifier::new("b")),
            Some(&Value::Int(20))
        ); // env2 shadows
        assert_eq!(
            env1.lookup(&SimpleIdentifier::new("c")),
            Some(&Value::Int(3))
        );
    }
}
