use crate::{
    error::{CodeRegion, Error, HasSpan},
    interpreter::Interpreter,
    profile::Profile,
    utils::get_instruction_attr,
};
use arcis_compiler::{
    compile::{write_ir, ArcisCOptions},
    profile_summary::ProfileSummary,
};
use arcis_interface::write_interface;
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use std::{
    collections::HashMap,
    fs::{create_dir_all, OpenOptions},
    io::Write,
};

#[derive(Debug, Default)]
pub struct CircuitBuilder<'ast> {
    interpreter: Interpreter<'ast>,
    pub tests: TokenStream,
    pub errors: Vec<Error>,
    pub n_circuits: usize,
    mod_depth: usize,
    pre_profiles: HashMap<String, Profile>,
    circuit_locations: HashMap<String, CodeRegion>,
}

impl<'ast> CircuitBuilder<'ast> {
    pub fn new(interpreter: Interpreter<'ast>) -> Self {
        CircuitBuilder {
            interpreter,
            ..Default::default()
        }
    }

    pub fn build_circuit(
        &mut self,
        item_fn: &'ast syn::ItemFn,
        attr: &'ast syn::Attribute,
    ) -> Result<TokenStream, Error> {
        let result = self.interpreter.start_and_run(item_fn)?;
        let (ir, test_code, inputs_ident, ret_ident, interface, pre_profile) = result.destructure();
        let options = ArcisCOptions { out_dir: None };
        let circuit_name = interface.name.clone();
        self.circuit_locations
            .insert(circuit_name.clone(), attr.span().region());
        let ir_path = write_ir(ir, &circuit_name, &options).unwrap();
        write_interface(interface.serialize().unwrap(), &circuit_name, &None).unwrap();
        self.pre_profiles.insert(circuit_name.clone(), pre_profile);
        let test_name = format_ident!("generated_test_{}", item_fn.sig.ident);
        Ok(quote! {
            #[cfg(test)]
            #[cfg(debug_assertions)]
            #[test]
            fn #test_name() {
                use arcis::testing::*;
                use arcis::arcis_type::*;
                let rng = &mut thread_rng();
                let ir = read_ir(#ir_path).expect("Deserialization failed. This may be caused by stale build files, please try cleaning the build/ directory and rebuilding.");
                let n_test_rounds = (16384 / ir.get_exprs().len()).clamp(1, 16);
                for _ in 0..n_test_rounds {
                    let _ = generated_bools();
                    #test_code
                    let rng_bools = generated_bools();
                    let mock_res = ir.eval_vec(#inputs_ident.clone(), rng, &rng_bools).unwrap();
                    assert!(#ret_ident.is_similar(ArcisType::from_values(&mock_res)), "Test failed for inputs {:?}.\nThe rust code and the circuit did not give the same result.\nThis could be a result of integer overflow: make sure that overflow-checks are enabled.", #inputs_ident);
                }
                if std::env::var("TEST_CIRCUIT").is_ok() {
                    let output_domains = ir.get_output_domains();
                    let opt_ir = ir.clone().optimize_into_circuitable();
                    let mask = ir.generated_bools_filter(&opt_ir);
                    let circ = opt_ir.to_async_mpc_circuit();

                    for _ in 0..n_test_rounds {
                        let _ = generated_bools();
                        #test_code
                        let gen_bools = generated_bools();
                        let rng_bools: Vec<Option<bool>> = mask.iter().map(|a| a.and_then(|idx| gen_bools.get(idx).copied())).collect();
                        let mock_res = circ.mock_eval_vec(#inputs_ident.clone(), &output_domains, &rng_bools, rng);
                        assert!(#ret_ident.is_similar(ArcisType::from_values(&mock_res)), "mock circuit test failed for inputs {:?}.", #inputs_ident);
                    }
                    if arcis::performance::is_performance_tracking_enabled() {
                        arcis::performance::track_instruction_performance(#circuit_name, circ.profile_info()).expect("Performance issue:");
                    }
                }
            }
        })
    }
    pub fn write_post_build(&self, mod_name: String) -> Result<(), std::io::Error> {
        let current_dir = std::env::current_dir()?;
        let circuits_dir = format!(
            "{}/build",
            current_dir.to_str().expect("current dir not a string")
        );
        create_dir_all(&circuits_dir)?;

        // pre_profile
        {
            let file_path = format!("{}/{}.pre_profile.js", circuits_dir, mod_name);
            let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&file_path)?;
            file.write_all(ProfileSummary::weight_js().as_bytes())?;
            file.write_all("var pre_profile = { 'code':\n".as_bytes())?;
            file.write_all(
                serde_json::to_string(&self.interpreter.code_store)
                    .expect("Serialization failed.")
                    .as_bytes(),
            )?;
            file.write_all("\n, 'instructions':\n".as_bytes())?;

            file.write_all(
                serde_json::to_string(&self.pre_profiles)
                    .expect("Serialization failed.")
                    .as_bytes(),
            )?;
            file.write_all("\n, 'instruction_locations':\n".as_bytes())?;

            file.write_all(
                serde_json::to_string(&self.circuit_locations)
                    .expect("Serialization failed.")
                    .as_bytes(),
            )?;
            file.write_all("\n};\n".as_bytes())?;
        }
        // client_types
        {
            let file_path = format!("{}/{}.client_types.json", circuits_dir, mod_name);
            let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&file_path)?;
            file.write_all(self.interpreter.client_types.to_json().as_bytes())?;
        }
        // client_types TypeScript
        {
            let file_path = format!("{}/{}.ts", circuits_dir, mod_name);
            let mut file = OpenOptions::new()
                .write(true)
                .truncate(true)
                .create(true)
                .open(&file_path)?;
            file.write_all(
                self.interpreter
                    .client_types
                    .to_typescript(&mod_name)
                    .as_bytes(),
            )?;
        }
        Ok(())
    }
}

impl<'ast> syn::visit::Visit<'ast> for CircuitBuilder<'ast> {
    fn visit_item_fn(&mut self, i: &'ast syn::ItemFn) {
        let is_circuit = get_instruction_attr(i);
        if let Some(attr) = is_circuit {
            if self.mod_depth == 0 {
                self.n_circuits += 1;
                let res = self.build_circuit(i, attr);
                match res {
                    Ok(circuit) => self.tests.extend(circuit),
                    Err(err) => self.errors.push(err),
                }
            } else {
                self.errors.push(Error::new_spanned(
                    attr,
                    "Instructions are only allowed in your #[encrypted] module's top level.",
                ))
            }
        }
        syn::visit::visit_item_fn(self, i);
    }
    fn visit_item_mod(&mut self, i: &'ast syn::ItemMod) {
        self.mod_depth += 1;
        syn::visit::visit_item_mod(self, i);
        self.mod_depth -= 1;
    }
}
