eprocedural 0.1.1

use std::collections::HashMap;
use std::sync::{Arc, LazyLock, Mutex};

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::{braced, parse_macro_input, Ident, Token};
use syn::parse::{Parse, ParseStream, Result};

struct ProcMacroInput {
    block_name: Ident,
    code: proc_macro2::TokenStream,
    proc_data: ProcedureData,
}

impl Parse for ProcMacroInput {
    fn parse(input: ParseStream) -> Result<Self> {
        let block_name: Ident = input.parse()?;

        let mut with_clause: Vec<(Ident, Ident)> = vec![];

        match input.parse::<Ident>() {
            Ok(ident) => {
                if ident == "with" {
                    loop {
                        let Ok(name) = input.parse::<Ident>() else { break };
                        input.parse::<Token![:]>()?;
                        let type_name = input.parse::<Ident>()?;
                        with_clause.push((name, type_name));

                        let _ = input.parse::<Token![,]>();
                    }
                }
            }
            Err(_) => {}
        }

        let content;
        braced!(content in input);
        let code = content.parse::<proc_macro2::TokenStream>()?;

        let mut defines = vec![];
        if let Ok(def) = input.parse::<Ident>() {
            if def != "defines" {
                return Err(syn::Error::new(def.span(), format!("unexpected identifier {}, expected 'defines'", def.to_string())));
            }

            loop {
                let Ok(name) = input.parse::<Ident>() else { break };
                input.parse::<Token![:]>()?;
                let type_name = input.parse::<Ident>()?;
                defines.push(ProcedureDefine {
                    name: name.to_string(),
                    type_name: type_name.to_string()
                });
    
                let _ = input.parse::<Token![,]>();
            }
        }

        let proc_data = ProcedureData {
            defines,
            imports: with_clause.iter().map(|m| ProcedureImport {
                name: m.0.to_string(),
                type_name: m.1.to_string(),
            }).collect()
        };

        Ok(Self {
            block_name,
            code,
            proc_data,
        })
    }
}

#[derive(Debug)]
struct ProcedureData {
    imports: Vec<ProcedureImport>,
    defines: Vec<ProcedureDefine>,
}

#[derive(Debug)]
struct ProcedureImport {
    name: String,
    type_name: String,
}

#[derive(Debug)]
struct ProcedureDefine {
    name: String,
    type_name: String,
}

static IDENTS_TO_THEIR_DEFINES: LazyLock<Arc<Mutex<HashMap<String, ProcedureData>>>> = std::sync::LazyLock::new(|| Arc::new(Mutex::new(HashMap::new())));

#[proc_macro]
pub fn proc(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ProcMacroInput);

    let block_name = input.block_name;
    let code = input.code;
    let defines_list: Vec<Ident> = input.proc_data.defines.iter().map(|m| Ident::new(&m.name, Span::call_site())).collect();

    let defines_types: Vec<Ident> = input.proc_data.defines.iter().map(|m| Ident::new(&m.type_name, Span::call_site())).collect();

    let imports_names_types: Vec<NameAndType> = input.proc_data.imports
        .iter()
        .map(|n| NameAndType {
                name: Ident::new(&n.name, Span::call_site()),
                type_: Ident::new(&n.type_name, Span::call_site())
            }
        )
        .collect();
    
    let mut data = IDENTS_TO_THEIR_DEFINES.lock().unwrap();
    data.insert(block_name.to_string(), input.proc_data);

    quote! {
        #[allow(non_camel_case_types)]
        struct #block_name {}

        impl #block_name {
            pub fn call(#(#imports_names_types),*) -> (#(#defines_types),*) {
                #code

                return (#(#defines_list),*)
            }
        }
    }.into()
}

struct NameAndType {
    name: Ident,
    type_: Ident,
}

impl quote::ToTokens for NameAndType {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        let name = self.name.clone();
        let type_ = self.type_.clone();
        tokens.extend(quote! {
            #name : &mut #type_
        });
    }
}

#[proc_macro]
pub fn with(input: TokenStream) -> TokenStream {
    let ident = parse_macro_input!(input as Ident).to_string();
    
    let map = IDENTS_TO_THEIR_DEFINES.lock().unwrap();
    let proc_data = map.get(&ident).unwrap();
    let defines_data: Vec<Ident> = proc_data.defines.iter().map(|m| Ident::new(&m.name, Span::call_site())).collect();
    let imports_data: Vec<Ident> = proc_data.imports.iter().map(|m| Ident::new(&m.name, Span::call_site())).collect();

    let ident = Ident::new(&ident, Span::call_site());

    return quote! {
        let ( #(mut #defines_data),* ) = #ident::call(#(&mut #imports_data),*);
    }.into();
}