use core::fmt::Debug;
use std::rc::Rc;

use hashbrown::HashMap;

use crate::error::Result;
use crate::ty::internal::Value;

#[derive(Clone, Debug)]
pub struct Env {
    cache: Vec<HashMap<usize, Value>>,
    with: Vec<Rc<HashMap<String, Value>>>,
    args: Vec<Value>,
}

impl Env {
    pub fn new() -> Self {
        Self {
            cache: Vec::from([HashMap::new()]),
            with: Vec::new(),
            args: Vec::new(),
        }
    }

    pub fn with_new_cache<T>(
        &mut self,
        f: impl FnOnce(&mut Self) -> T,
    ) -> (T, HashMap<usize, Value>) {
        self.cache.push(HashMap::new());
        let ret = f(self);
        (ret, self.cache.pop().unwrap())
    }

    pub fn with_cache<T>(
        &mut self,
        cache: HashMap<usize, Value>,
        f: impl FnOnce(&mut Self) -> T,
    ) -> (T, HashMap<usize, Value>) {
        self.cache.push(cache);
        let ret = f(self);
        (ret, self.cache.pop().unwrap())
    }

    pub fn insert_cache(&mut self, idx: usize, val: Value) {
        self.cache.last_mut().unwrap().insert(idx, val);
    }

    pub fn insert_cache_lazy(
        &mut self,
        idx: usize,
        f: impl FnOnce(&mut Self) -> Result<Value>,
    ) -> Result<()> {
        if self.cache.last().unwrap().get(&idx).is_none() {
            let val = f(self)?;
            self.cache.last_mut().unwrap().insert(idx, val);
        }
        Ok(())
    }

    pub fn lookup_cache(
        &mut self,
        idx: usize,
        f: impl FnOnce(&mut Env) -> Result<Value>,
    ) -> Result<Value> {
        for level in self.cache.iter().rev() {
            if let Some(ret) = level.get(&idx) {
                return ret.clone().ok();
            }
        }
        let val = f(self)?;
        self.cache.last_mut().unwrap().insert(idx, val.clone());
        val.ok()
    }

    pub fn lookup_arg(&self, level: usize) -> Value {
        self.args[self.args.len() - level - 1].clone()
    }

    pub fn lookup_with(&self, symbol: &str) -> Option<Value> {
        for with in self.with.iter().rev() {
            if let Some(ret) = with.get(symbol) {
                return Some(ret.clone());
            }
        }
        None
    }

    pub fn enter_arg(&mut self, arg: Value) {
        self.args.push(arg);
    }

    pub fn pop_args(&mut self, len: usize) -> Vec<Value> {
        self.args.split_off(self.args.len() - len)
    }

    pub fn reserve_args(&mut self, len: usize) {
        self.args.reserve(len);
    }

    pub fn enter_args(&mut self, args: Vec<Value>) {
        self.args.extend(args);
    }

    pub fn pop_with(&mut self) {
        self.with.pop();
    }

    pub fn enter_with(&mut self, map: Rc<HashMap<String, Value>>) {
        self.with.push(map)
    }
}