use hashbrown::HashMap;
use std::rc::Rc;

use crate::ty::internal::{AttrSet, Value};

#[derive(Debug, Default, Clone)]
pub struct LetEnv<'jit, 'vm> {
    map: Rc<HashMap<usize, Value<'jit, 'vm>>>,
    last: Option<Rc<LetEnv<'jit, 'vm>>>,
}

#[derive(Debug, Default, Clone)]
pub struct WithEnv<'jit, 'vm> {
    map: Rc<AttrSet<'jit, 'vm>>,
    last: Option<Rc<WithEnv<'jit, 'vm>>>,
}

impl<'jit, 'vm> LetEnv<'jit, 'vm> {
    pub fn new(map: Rc<HashMap<usize, Value<'jit, 'vm>>>) -> Self {
        Self { map, last: None }
    }

    pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
        if let Some(val) = self.map.get(&symbol).cloned() {
            return Some(val);
        }
        self.last.as_ref().map(|env| env.lookup(symbol)).flatten()
    }

    pub fn enter_let(self, new: impl Iterator<Item = (usize, Value<'jit, 'vm>)>) -> Self {
        let map = Rc::new(new.collect());
        let last = Some(self.into());
        LetEnv { last, map }
    }

    pub fn enter_with(self, new: Rc<AttrSet<'jit, 'vm>>) -> Self {
        let map = new
            .as_inner()
            .iter()
            .map(|(&k, v)| {
                (
                    k,
                    if let Value::Builtins(weak) = v {
                        Value::AttrSet(weak.upgrade().unwrap())
                    } else {
                        v.clone()
                    },
                )
            })
            .collect::<HashMap<_, _>>()
            .into();
        let last = Some(self.into());
        LetEnv { last, map }
    }

    pub fn leave(self) -> Self {
        self.last.unwrap().as_ref().clone()
    }
}

impl<'jit, 'vm> WithEnv<'jit, 'vm> {
    pub fn lookup(&self, symbol: usize) -> Option<Value<'jit, 'vm>> {
        if let Some(val) = self.map.select(symbol) {
            return Some(val);
        }
        self.last.as_ref().map(|env| env.lookup(symbol)).flatten()
    }

    pub fn enter_with(self, new: Rc<AttrSet<'jit, 'vm>>) -> Self {
        let map = Rc::new(new
            .as_inner()
            .iter()
            .map(|(&k, v)| {
                (
                    k,
                    if let Value::Builtins(weak) = v {
                        Value::AttrSet(weak.upgrade().unwrap())
                    } else {
                        v.clone()
                    },
                )
            })
            .collect::<HashMap<_, _>>()
            .into());
        let last = Some(self.into());
        WithEnv { last, map }
    }

    pub fn leave(self) -> Self {
        self.last.unwrap().as_ref().clone()
    }
}