use std::iter::IntoIterator;

use proc_macro;
use syn::parse_macro_input;

fn tuple_type_from_sig<'a, Iter: IntoIterator<Item = &'a syn::FnArg>>(
    args: Iter,
) -> syn::TypeTuple {
    let types = args
        .into_iter()
        .map(|a| match &a {
            &syn::FnArg::Receiver(_) => panic!("unexpected self in function argument list"),
            &syn::FnArg::Typed(ref typ) => typ.ty.clone(),
        })
        .collect::<Vec<Box<syn::Type>>>();

    if types.len() == 1 {
        let mut types = types;
        let typ = types.pop().unwrap();
        syn::parse_quote! { (#typ,) }
    } else {
        syn::parse_quote! {
            (#(#types),*)
        }
    }
}

#[proc_macro_attribute]
pub fn lazy(_: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let parsedfn = parse_macro_input!(item as syn::ItemFn);
    let sig = &parsedfn.sig;
    let fnname = &sig.ident;
    let returntype = match &sig.output {
        syn::ReturnType::Default => {
            let emptytuplets = proc_macro::TokenStream::from(quote::quote! { () });
            let emptytuple: syn::Type = parse_macro_input!(emptytuplets as syn::Type);
            syn::ReturnType::Type(<syn::Token![->]>::default(), Box::new(emptytuple))
        }
        r => r.clone(),
    };

    let newname = syn::Ident::new((fnname.to_string() + "_lazy").as_str(), fnname.span());
    let mut newsig = sig.clone();
    newsig.ident = newname.clone();

    match sig.inputs.first() {
        Some(syn::FnArg::Receiver(_)) => {
            // this is an impl function

            // Skip `self` argument.
            let mut args_iter = (&newsig.inputs).into_iter();
            let receiver = args_iter.next().unwrap();
            let tuty = tuple_type_from_sig(args_iter);
            // We are not evaluating for the receiver, thus one fewer argument.
            let args_ixs: Vec<usize> = (0..newsig.inputs.len()-1).collect();
            (quote::quote! {
                #parsedfn

                fn #newname<'lazyvallt>(#receiver, lzy: ::lazycall_support::LazyVal<'lazyvallt, #tuty>) #returntype {
                    use ::lazycall_support::Evaluate;
                    let evald = lzy.eval();
                    self.#fnname(#(evald.#args_ixs),*)
                }
            }).into()
        }
        Some(syn::FnArg::Typed(_)) => {
            // This is a freestanding function
            let tuty = tuple_type_from_sig(&newsig.inputs);
            let args_ixs: Vec<usize> = (0..newsig.inputs.len()).collect();
            (quote::quote! {
                #parsedfn

                fn #newname<'lazyvallt>(lzy: ::lazycall_support::LazyVal<'lazyvallt, #tuty>) #returntype {
                    use ::lazycall_support::Evaluate;
                    let evald = lzy.eval();
                    #fnname(#(evald.#args_ixs),*)
                }
            }).into()
        }
        None => {
            // This is a freestanding function without arguments.
            // fn f() -> ResultT { ... } => fn f_lazy() { ... } (no body modification)
            let mut updatedfn = parsedfn.clone();
            updatedfn.sig = newsig;
            (quote::quote! {
                #updatedfn

                #parsedfn
            })
            .into()
        }
    }
}