use crate::utils::{get_base_circuit, ARCIS_CIRCUIT};
use quote::{quote, ToTokens};

pub struct HelperCleaner;

impl HelperCleaner {
    pub fn clean(module: &mut syn::ItemMod) {
        syn::visit_mut::visit_item_mod_mut(&mut HelperCleaner, module);
    }
}

fn into_arcis_circuit(attrs: &[syn::Attribute], block: &mut syn::Block, sig: &syn::Signature) {
    if let Some(circ) = get_base_circuit(attrs).cloned() {
        let inputs = sig.inputs.iter().map(|x| match x {
            syn::FnArg::Receiver(_) => x.to_token_stream(),
            syn::FnArg::Typed(a) => a.pat.to_token_stream(),
        });

        let code = quote! {
            {
                let mut inputs = Vec::new();
                arcis_type::ArcisType::handle_outputs(&(#(#inputs),*), &mut inputs);
                let circ = ArcisCircuit::new_from_str(#circ).expect("Not found");
                let outputs = circ.apply(inputs);
                arcis_type::ArcisType::from_values(&outputs)
            }
        };

        let new_block = syn::parse2(code).unwrap();
        *block = new_block;
    }
}

impl syn::visit_mut::VisitMut for HelperCleaner {
    fn visit_item_struct_mut(&mut self, i: &mut syn::ItemStruct) {
        let meta: syn::Meta = syn::parse2(quote! { derive(ArcisType) }).unwrap();
        i.attrs.push(syn::Attribute {
            pound_token: Default::default(),
            style: syn::AttrStyle::Outer,
            bracket_token: Default::default(),
            meta,
        });
        syn::visit_mut::visit_item_struct_mut(self, i);
    }
    fn visit_attributes_mut(&mut self, i: &mut Vec<syn::Attribute>) {
        i.retain_mut(|x| {
            let Some(ident) = x.meta.path().get_ident().cloned() else {
                return true;
            };
            if ident == "instruction" {
                x.meta = syn::parse2(quote! { allow(dead_code) })
                    .expect("`allow(dead_code)` is valid attribute, this should not fail.");
            }
            ident != ARCIS_CIRCUIT
        });
    }
    fn visit_item_fn_mut(&mut self, i: &mut syn::ItemFn) {
        into_arcis_circuit(&i.attrs, &mut i.block, &i.sig);
        syn::visit_mut::visit_item_fn_mut(self, i);
    }
    fn visit_impl_item_fn_mut(&mut self, i: &mut syn::ImplItemFn) {
        into_arcis_circuit(&i.attrs, &mut i.block, &i.sig);
        syn::visit_mut::visit_impl_item_fn_mut(self, i);
    }
    fn visit_item_use_mut(&mut self, i: &mut syn::ItemUse) {
        if i.tree == syn::parse_quote! {arcis::*} || i.tree == syn::parse_quote! {arcis_imports::*}
        {
            i.attrs.push(syn::parse_quote! {#[allow(unused_imports)]});
        }
        syn::visit_mut::visit_item_use_mut(self, i);
    }
}
