use proc_macro2::TokenStream;
use quote::{format_ident, quote};

pub struct TestCode {
    fn_to_produce_inputs: TokenStream,
    input_idents: Vec<syn::Ident>,
    input_values_ident: syn::Ident,
    ident_id: usize,
}

impl TestCode {
    pub fn new() -> Self {
        let input_values_ident = format_ident!("input_values");
        let fn_to_produce_inputs = quote! {
            let mut #input_values_ident: Vec<EvalValue> = Vec::new();
        };
        TestCode {
            fn_to_produce_inputs,
            input_idents: Vec::new(),
            input_values_ident,
            ident_id: 0,
        }
    }
    fn new_ident(&mut self) -> syn::Ident {
        let res = format_ident!("val_{}", self.ident_id);
        self.ident_id += 1;
        res
    }

    pub fn handle_input(&mut self) {
        let ident = self.new_ident();
        self.input_idents.push(ident.clone());
        let input_values_ident = &self.input_values_ident;
        self.fn_to_produce_inputs.extend(quote! {
            let #ident = ArcisType::gen_input(&mut #input_values_ident);
        });
    }
    pub fn finish(mut self, circuit_name: &syn::Ident) -> (TokenStream, syn::Ident, syn::Ident) {
        let return_value_ident = format_ident!("{}_return", circuit_name);
        let input_idents = self.input_idents.iter();
        self.fn_to_produce_inputs.extend(quote! {
            let #return_value_ident = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| #circuit_name(#(#input_idents),*)));
            let Ok(#return_value_ident) = #return_value_ident else {continue;};
        });
        (
            self.fn_to_produce_inputs,
            self.input_values_ident,
            return_value_ident,
        )
    }
}
