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

use crate::ty::common::Const;
use crate::ty::internal::{AttrSet, PrimOp, Value};
use crate::vm::{VM, VmEnv};

pub fn env<'jit, 'vm>(vm: &'vm VM<'jit>) -> VmEnv<'jit, 'vm> {
    let mut env_map = HashMap::new();
    env_map.insert(vm.new_sym("true"), Value::Const(Const::Bool(true)));
    env_map.insert(vm.new_sym("false"), Value::Const(Const::Bool(false)));

    let primops = [
        PrimOp::new("add", 2, |_, args| {
            let [first, second]: [Value; 2] = args.try_into().unwrap();
            first.add(second).ok()
        }),
        PrimOp::new("sub", 2, |_, args| {
            let [first, second]: [Value; 2] = args.try_into().unwrap();
            first.add(second.neg()).ok()
        }),
        PrimOp::new("mul", 2, |_, args| {
            let [first, second]: [Value; 2] = args.try_into().unwrap();
            first.mul(second).ok()
        }),
        PrimOp::new("div", 2, |_, args| {
            let [first, second]: [Value; 2] = args.try_into().unwrap();
            first.div(second)
        }),
        PrimOp::new("lessThan", 2, |_, args| {
            let [first, second]: [Value; 2] = args.try_into().unwrap();
            first.lt(second).ok()
        }),
        PrimOp::new("seq", 2, |vm, args| {
            let [mut first, second]: [Value; 2] = args.try_into().unwrap();
            first.force(vm).unwrap();
            second.ok()
        }),
        PrimOp::new("deepSeq", 2, |vm, args| {
            let [mut first, second]: [Value; 2] = args.try_into().unwrap();
            first.force_deep(vm).unwrap();
            second.ok()
        }),
    ];

    let mut map = HashMap::new();
    for primop in primops {
        let primop = Rc::new(primop);
        env_map.insert(
            vm.new_sym(format!("__{}", primop.name)),
            Value::PrimOp(primop.clone()),
        );
        map.insert(vm.new_sym(primop.name), Value::PrimOp(primop));
    }
    let sym = vm.new_sym("builtins");
    let attrs = Rc::new_cyclic(|weak| {
        map.insert(sym, Value::Builtins(weak.clone()));
        AttrSet::from_inner(map)
    });
    let builtins = Value::AttrSet(attrs);

    env_map.insert(sym, builtins);
    VmEnv::new(env_map.into())
}