use crate::core::ir_builder::IRBuilder;
use std::cell::Cell;

thread_local! {
    static EXPR_STORE: Cell<IRBuilder> = Cell::new(IRBuilder::new(true));
}

fn swap_expr_store(expr_store: &mut IRBuilder) {
    let local_cell = Cell::from_mut(expr_store);
    EXPR_STORE.with(|cell| cell.swap(local_cell));
}

#[allow(dead_code)]
pub fn with_local_expr_store_as_global<FN, R>(f: FN, expr_store: &mut IRBuilder) -> R
where
    FN: FnOnce() -> R,
{
    swap_expr_store(expr_store);
    let res = f();
    swap_expr_store(expr_store);
    res
}

pub fn with_global_expr_store_as_local<F, R>(f: F) -> R
where
    F: FnOnce(&mut IRBuilder) -> R,
{
    let mut temp = IRBuilder::new(true);
    swap_expr_store(&mut temp);
    let res = f(&mut temp);
    swap_expr_store(&mut temp);
    assert_eq!(temp.len(), 0, "Problem with expr_store, you are using global expr_store inside of a with_global_expr_store_as_local.");
    res
}

#[allow(dead_code)] // The function is actually used, but Clippy doesn't realize.
pub fn take_global_expr_store() -> IRBuilder {
    let mut temp = IRBuilder::new(true);
    swap_expr_store(&mut temp);
    temp
}

pub fn are_unbuilt_circuits_allowed() -> bool {
    with_global_expr_store_as_local(|expr_store| expr_store.are_unbuilt_circuits_allowed())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::{
            bounds::{Bounds, FieldBounds, IsBounds},
            expressions::{
                expr::Expr,
                field_expr::{FieldExpr, InputInfo},
                InputKind,
            },
            global_value::value::FieldValue,
        },
        utils::field::ScalarField,
    };
    use rustc_hash::FxHashMap;
    use std::rc::Rc;

    #[test]
    fn test_chain() {
        let rng = &mut crate::utils::test_rng::get();
        let mut expr_store = IRBuilder::new(true);
        let input = expr_store.new_expr(Expr::Scalar(FieldExpr::Input(
            0,
            Rc::new(InputInfo {
                kind: InputKind::Secret,
                min: 0.into(),
                max: 7.into(),
                ..InputInfo::default()
            }),
        )));
        let res = with_local_expr_store_as_global(
            || {
                let sum_8 = FieldValue::<ScalarField>::from_id(input) + FieldValue::from(8);
                with_global_expr_store_as_local(|expr_store| {
                    let val_16 = expr_store.new_expr(Expr::Scalar(FieldExpr::Val(16.into())));
                    let sum_16 =
                        expr_store.new_expr(Expr::Scalar(FieldExpr::Add(sum_8.get_id(), val_16)));
                    with_local_expr_store_as_global(
                        || {
                            let sum_32 =
                                FieldValue::<ScalarField>::from_id(sum_16) + FieldValue::from(32);
                            with_global_expr_store_as_local(|expr_store| {
                                let val_64 =
                                    expr_store.new_expr(Expr::Scalar(FieldExpr::Val(64.into())));
                                expr_store
                                    .new_expr(Expr::Scalar(FieldExpr::Add(sum_32.get_id(), val_64)))
                            })
                        },
                        expr_store,
                    )
                })
            },
            &mut expr_store,
        );
        let ir = expr_store.into_ir(vec![res]);
        let mut hash_map: FxHashMap<usize, _> = FxHashMap::default();
        let eval_res = ir.eval(rng, &mut hash_map);
        assert!(eval_res.is_ok());
        let eval_res = eval_res.unwrap();
        let expected_res_bounds = Bounds::Scalar(FieldBounds::new(120.into(), 127.into()));
        assert!(expected_res_bounds.contains(eval_res[0]))
    }
}
