use crate::core::{
    circuits::boolean::boolean_value::BooleanValue,
    expressions::expr::EvalFailure,
    global_value::global_expr_store::with_local_expr_store_as_global,
    ir_builder::IRBuilder,
};
use std::fmt::Debug;

/// A trait to define a boolean sub-circuit.
pub trait BooleanCircuit: Debug {
    #[allow(dead_code)]
    /// The operation that is being performed.
    fn eval(&self, x: Vec<bool>) -> Result<Vec<bool>, EvalFailure>;

    #[allow(dead_code)]
    /// The operation, in MPC.
    fn run(&self, vals: Vec<BooleanValue>) -> Vec<BooleanValue>;

    #[allow(dead_code)]
    fn run_usize(&self, vals: &[usize], expr_store: &mut IRBuilder) -> Vec<usize> {
        with_local_expr_store_as_global(
            || self.run(vals.iter().map(|id| BooleanValue::new(*id)).collect()),
            expr_store,
        )
        .iter()
        .map(BooleanValue::get_id)
        .collect()
    }
}

#[cfg(test)]
pub mod tests {
    use super::*;
    use crate::{
        core::{
            circuits::traits::SAVE_CIRC_TEST_FOLDER_ENV_VAR,
            expressions::{
                bit_expr::{BitExpr, BitInputInfo},
                domain::Domain,
                expr::EvalValue,
                InputKind,
            },
            ir_builder::ExprStore,
        },
        utils::field::BaseField,
    };
    use rand::Rng;
    use std::rc::Rc;

    fn desc_file_path() -> String {
        let folder = std::env::var(SAVE_CIRC_TEST_FOLDER_ENV_VAR).unwrap();
        let binding = std::thread::current();
        let test_name = binding.name().unwrap();
        format!("{folder}/{}.desc", test_name)
    }
    fn run_file_path() -> String {
        let folder = std::env::var(SAVE_CIRC_TEST_FOLDER_ENV_VAR).unwrap();
        let binding = std::thread::current();
        let test_name = binding.name().unwrap();
        format!("{folder}/{}.run", test_name)
    }

    fn test<R: Rng + ?Sized, C: TestedBooleanCircuit>(rng: &mut R, desc: &C) {
        let n_inputs = desc.gen_n_inputs(rng);
        let input_vals = (0..n_inputs)
            .map(|_| rng.gen_bool(0.5))
            .collect::<Vec<bool>>();
        let eval_result = desc.eval(input_vals.clone());
        // This circuit will be built by the compiler when circuits should be built.
        let mut expr_store = IRBuilder::new(false);
        let input_ids = (0..n_inputs)
            .map(|i| {
                <IRBuilder as ExprStore<BaseField>>::push_bit(
                    &mut expr_store,
                    BitExpr::Input(
                        i,
                        Rc::new(BitInputInfo {
                            kind: InputKind::Secret,
                            ..BitInputInfo::default()
                        }),
                    ),
                )
            })
            .collect::<Vec<usize>>();
        let outputs = desc.run_usize(&input_ids, &mut expr_store);
        let test_ir = expr_store.into_ir(outputs);
        let mut input_vals_map = input_vals
            .iter()
            .map(|b| EvalValue::Bit(*b))
            .enumerate()
            .collect();
        let test_result = test_ir
            .eval(rng, &mut input_vals_map)
            .map(|x| x.into_iter().map(bool::unwrap).collect::<Vec<bool>>())
            .unwrap();
        if eval_result.is_err() {
            return;
        }
        let eval_result = eval_result.unwrap();
        assert_eq!(
            test_result, eval_result,
            "different results in circuit\n{:?}\nfor inputs\n{:?}\n",
            desc, input_vals
        );
        desc.extra_checks(input_vals, eval_result)
    }

    /// A trait to test BooleanCircuits.
    pub trait TestedBooleanCircuit: BooleanCircuit + Clone + 'static {
        /// A function to randomly generate a description of the sub-circuit.
        fn gen_desc<R: Rng + ?Sized>(rng: &mut R) -> Self;

        /// A function to randomly generate the number of inputs for the sub-circuit
        fn gen_n_inputs<R: Rng + ?Sized>(&self, rng: &mut R) -> usize;

        /// A function that can be overwritten in impls to perform extra-checks.
        #[allow(unused_variables)]
        fn extra_checks(&self, inputs: Vec<bool>, outputs: Vec<bool>) {}

        /// The actual test.
        /// n_desc is the number of different descriptions that will be generated.
        /// n_runs is the number of runs that will be attempted per description.
        /// A run where eval fails is a successful try.
        fn test(n_desc: usize, n_runs_per_desc: usize) {
            let rng = &mut crate::utils::test_rng::get();
            let (save_desc, save_run) = if std::env::var(SAVE_CIRC_TEST_FOLDER_ENV_VAR).is_ok() {
                let desc_path = desc_file_path();
                println!("saving the circuit description at {}", desc_path);
                let run_path = run_file_path();
                println!("saving the circuit run at {}", run_path);
                (Some(desc_path), Some(run_path))
            } else {
                (None, None)
            };
            for _ in 0..n_desc {
                if let Some(file_path) = &save_desc {
                    crate::utils::test_rng::save_to_file(rng, file_path);
                }
                let desc = Self::gen_desc(rng);
                for _ in 0..n_runs_per_desc {
                    if let Some(file_path) = &save_run {
                        crate::utils::test_rng::save_to_file(rng, file_path);
                    }
                    test(rng, &desc);
                }
            }
        }
    }
}
