use std::{cell::RefCell, collections::HashMap, rc::Rc};

use thiserror::Error;

use crate::{token::Token, value::Value};

#[derive(Error, Debug)]
pub enum EnvironmentError {
    #[error("line {0}: undefined variable: {1}")]
    UndefinedVariable(usize, String),
    #[error("invalid environment distance")]
    InvalidDistance,
}

/// Environment mapping variable names to their respective values. Can have an optional enclosing
/// environment. The outermost (global) environment has no enclosing one.
#[derive(Default, Debug, Clone)]
pub struct Environment {
    values: HashMap<String, Value>,
    enclosing: Option<Rc<RefCell<Environment>>>,
}

impl Environment {
    /// Initialize a new environment with an enclosing one.
    pub fn with_enclosing(enclosing: Rc<RefCell<Environment>>) -> Self {
        Self {
            values: HashMap::default(),
            enclosing: Some(enclosing),
        }
    }

    /// Define a new variable binding with a value.
    ///
    /// Overwrite any binding of the same name.
    pub fn define(&mut self, name: String, value: Value) {
        self.values.insert(name, value);
    }

    /// Assign a new value to an existing (defined) variable. Error if there is no such variable in
    /// this environment or any of the enclosing ones.
    pub fn assign(&mut self, token: &Token, value: Value) -> Result<(), EnvironmentError> {
        if self.values.contains_key(token.lexeme.as_str()) {
            self.values.insert(token.lexeme.clone(), value);
            Ok(())
        } else if let Some(enclosing) = &mut self.enclosing {
            enclosing.borrow_mut().assign(token, value)
        } else {
            Err(EnvironmentError::UndefinedVariable(
                token.line,
                token.lexeme.clone(),
            ))
        }
    }

    pub fn assign_at(
        &mut self,
        distance: usize,
        token: &Token,
        value: Value,
    ) -> Result<(), EnvironmentError> {
        self.ancestor(distance)?
            .values
            .insert(token.lexeme.clone(), value);
        Ok(())
    }

    /// Get the value of an existing (defined) variable. Error if there is no such variable in
    /// this environment or any of the enclosing ones.
    pub fn get(&self, token: &Token) -> Result<Value, EnvironmentError> {
        if let Some(v) = self.values.get(token.lexeme.as_str()) {
            Ok(v.clone())
        } else if let Some(enclosing) = &self.enclosing {
            enclosing.borrow().get(token)
        } else {
            Err(EnvironmentError::UndefinedVariable(
                token.line,
                token.lexeme.clone(),
            ))
        }
    }

    pub fn get_at(&self, distance: usize, token: &Token) -> Result<Value, EnvironmentError> {
        if let Some(v) = self.ancestor(distance)?.values.get(token.lexeme.as_str()) {
            Ok(v.clone())
        } else {
            Err(EnvironmentError::UndefinedVariable(
                token.line,
                token.lexeme.clone(),
            ))
        }
    }

    fn ancestor(&self, distance: usize) -> Result<Environment, EnvironmentError> {
        let mut environment = self.clone();

        for _ in 0..distance {
            let enclosing = match &environment.enclosing {
                Some(enc) => enc.borrow().clone(),
                None => return Err(EnvironmentError::InvalidDistance),
            };
            environment = enclosing;
        }

        Ok(environment)
    }
}