extern crate proc_macro;

use proc_macro::TokenStream;
use quote::quote;
use syn::{AttributeArgs, ItemFn, Lit, Meta, parse_macro_input, parse_quote};

#[proc_macro_attribute]
pub fn perf(attr: TokenStream, item: TokenStream) -> TokenStream {
    let attrs = parse_macro_input!(attr as AttributeArgs);
    let name = match &attrs[0] {
        syn::NestedMeta::Lit(Lit::Str(l)) => l.value(),
        _ => panic!("Invalid attribute"),
    };
    let counter_type = match &attrs[1] {
        syn::NestedMeta::Meta(Meta::Path(p)) => p,
        _ => panic!("Invalid attribute"),
    };
    let mut function = parse_macro_input!(item as ItemFn);
    function.block.stmts.insert(0, parse_quote! {
        let __perf_start = std::time::Instant::now();
    });

    let timer_end = parse_quote! {
        unsafe {
            #counter_type::write_perf(#name, __perf_start.elapsed().as_micros());
        }
    };

    if let Some(last_stmt) = function.block.stmts.last() {
        match last_stmt {
            syn::Stmt::Expr(_) => function.block.stmts.insert(function.block.stmts.len() - 1, timer_end),
            _ => function.block.stmts.push(timer_end),
        }
    } else {
        function.block.stmts.push(timer_end);
    }
    quote!(#function).into()
}
