athene_macro 0.3.5

mod attribute;
mod handler;

use crate::controller::{
    attribute::ControllerAttr,
    handler::{ArgumentType, FunctionArgument, Handler, HandlerOption},
};
use proc_macro2::{Ident, Span, TokenStream};
use quote::{quote, quote_spanned, ToTokens};
use syn::{
    spanned::Spanned, AttributeArgs, Error, GenericArgument, ItemImpl, PathArguments, Result, Type,
};

pub fn expand_controller(args: AttributeArgs, input: ItemImpl) -> Result<TokenStream> {
    let controller_attr = ControllerAttr::new(args, &input)?;
    let handlers = handler::parse_handlers(input)?;

    let mod_ident = Ident::new(
        &format!("macro_{}", &controller_attr.ident),
        Span::call_site(),
    );

    let controller_impl =
        attribute::gen_controller_trait_impl(&controller_attr, handlers.as_slice());
    let struct_impl = gen_struct_impl(controller_attr.ident, handlers)?;

    Ok(quote! {
        mod #mod_ident {
            use super::*;
            use athene::serde_urlencoded;
            use std::{str::FromStr,collections::HashMap};
            
            #struct_impl

            #controller_impl
        }
    })
}

fn gen_struct_impl(controller_ident: Ident, handlers: Vec<Handler>) -> Result<TokenStream> {
    let mut handler_tokens = TokenStream::new();
    for handler in handlers {
        if handler.option.needs_wrapper_fn() {
            gen_handler(&mut handler_tokens, handler)?;
        } else {
            handler.method.to_tokens(&mut handler_tokens);
        }
    }

    Ok(quote! {
        impl #controller_ident {
            #handler_tokens
        }
    })
}

fn gen_handler(handler_tokens: &mut TokenStream, handler: Handler) -> Result<()> {
    let opts = handler.option;
    let method_ident = handler.method.sig.ident.clone();
    let route_span = method_ident.span();
    let return_type = handler.return_type;

    let mut method = handler.method;
    let mut method_string = method_ident.to_string();
    method_string.push_str("_raw");

    method.attrs.push(syn::parse_quote! {#[inline]});
    method.sig.ident = Ident::new(method_string.as_str(), Span::mixed_site());
    method.to_tokens(handler_tokens);

    let inner_method_ident = method.sig.ident;

    let mut body_stream = TokenStream::new();
    (quote! {}).to_tokens(&mut body_stream);
    gen_cookie_load(&mut body_stream, &opts);
    gen_query_load(&mut body_stream, &opts);
    let mut call_params_ident = Vec::with_capacity(opts.fn_arguments.len());
    let async_call = !opts.sync_handler;
    for arg in opts.fn_arguments.into_iter() {
        arg.gen_parameter(&mut body_stream, &mut call_params_ident)?;
    }
    let inner_call = gen_call_to_inner(inner_method_ident, call_params_ident, async_call);

    let t = quote_spanned! {route_span=>
        #[allow(unused_mut)] // ====================
        async fn #method_ident(&self, mut req: Request) -> Result<#return_type, Error> {
            #body_stream
            Ok(#inner_call)
        }
    };

    t.to_tokens(handler_tokens);

    Ok(())
}

fn gen_cookie_load(stream: &mut TokenStream, opts: &HandlerOption) {
    if opts.parse_cookies {
        (quote! {
            req.parse_cookies();
        })
        .to_tokens(stream);
    }
}

fn gen_query_load(stream: &mut TokenStream, opts: &HandlerOption) {
    if opts.parse_query {
        (quote! { // ============================= 
        let mut query = req.uri().query().map(serde_urlencoded::from_str::<HashMap<String, String>>).transpose()?.unwrap_or_default();
        })
        .to_tokens(stream);
    }
}

fn gen_call_to_inner(
    inner_method_ident: Ident,
    idents: Vec<Ident>,
    async_call: bool,
) -> TokenStream {
    let mut call = TokenStream::new();

    (quote! {self.#inner_method_ident}).to_tokens(&mut call);

    gen_call_params(idents).to_tokens(&mut call);

    if async_call {
        (quote! {.await}).to_tokens(&mut call);
    }

    call
}

fn gen_call_params(idents: Vec<Ident>) -> TokenStream {
    let mut params = TokenStream::new();
    let paren = syn::token::Paren {
        span: Span::call_site(),
    };

    paren.surround(&mut params, |params| {
        let mut ident_iter = idents.into_iter();

        if let Some(first_param) = ident_iter.next() {
            (quote! {#first_param}).to_tokens(params);

            for i in ident_iter {
                (quote! {, #i}).to_tokens(params)
            }
        }
    });

    params
}

impl FunctionArgument {
    pub fn gen_parameter(
        mut self,
        stream: &mut TokenStream,
        call_ident: &mut Vec<Ident>,
    ) -> Result<()> {
        let optional = if let ArgumentType::Option(a) = self.a_type {
            self.typ = self
                .typ
                .and_then(|t| t.path.segments.into_iter().next())
                .map(|first| first.arguments)
                .and_then(|a| {
                    if let PathArguments::AngleBracketed(p_a) = a {
                        return p_a.args.into_iter().next().and_then(|g_a| {
                            if let GenericArgument::Type(Type::Path(type_path)) = g_a {
                                return Some(type_path);
                            }

                            None
                        });
                    }

                    None
                });
            match a.as_ref() {
                ArgumentType::SelfType | ArgumentType::Request | ArgumentType::Cookie => {
                    return Err(Error::new(
                        self.typ.as_ref().map(|t| t.span()).unwrap_or_else(Span::call_site),
                        "Optional parameters are only allowed for quey params, route params, or body param (Json or Form)",
                    ));
                }
                ArgumentType::Option(_) => {
                    return Err(Error::new(
                        self.typ
                            .as_ref()
                            .map(|t| t.span())
                            .unwrap_or_else(Span::call_site),
                        "Option within option are not allowed as handler parameters",
                    ));
                }
                _ => self.a_type = *a,
            }
            true
        } else {
            false
        };

        match &self.a_type {
            ArgumentType::SelfType => {
                return Ok(());
            }
            ArgumentType::Request => {
                call_ident.push(Ident::new("req", Span::call_site()));
                return Ok(());
            }
            _ => {}
        }

        let ident = Ident::new(self.name.as_str(), Span::call_site());
        match &self.a_type {
            ArgumentType::Query => self.gen_query(stream, optional),
            ArgumentType::Json => self.gen_json_param(stream, optional),
            ArgumentType::Form => self.gen_form_param(stream, optional),
            ArgumentType::Cookie => self.gen_cookie_param(stream),
            ArgumentType::Params { is_query_param, .. } => {
                if *is_query_param {
                    self.gen_query_param(stream, optional);
                } else {
                    self.gen_path_param(stream, optional);
                }
            }
            _ => {}
        }

        call_ident.push(ident);
        Ok(())
    }
    #[allow(dead_code)]
    #[allow(unused_imports)]
    fn gen_form_param(&self, stream: &mut TokenStream, optional: bool) {
        let id = Ident::new(self.name.as_str(), Span::call_site());
        let typ_raw = self.typ.as_ref().expect("This should not happens");

        let typ = self
            .typ
            .as_ref()
            .and_then(|t| t.path.segments.first())
            .map(|first| &first.arguments)
            .and_then(|a| {
                if let PathArguments::AngleBracketed(p_a) = a {
                    return p_a.args.first().and_then(|g_a| {
                        if let GenericArgument::Type(Type::Path(type_path)) = g_a {
                            return Some(type_path);
                        }

                        None
                    });
                }

                None
            })
            .expect("This should not happens");

        (quote! { // ====================== 
            let #id = if let Some(form) = req.uri().query().map(|query_str|serde_urlencoded::from_str::<#typ>(query_str).map_err(Error::SerdeUrlDe)) {
                Some(form.map(|x| Form(x)))
            } else {
                //// =========== \\\\\\
                Some(req.parse_form::<#typ_raw>().await)
            }.transpose()
        })
        .to_tokens(stream);

        if optional {
            (quote! {.ok().flatten();}).to_tokens(stream);
        } else {
            // ===============
            (quote! {?.ok_or_else(|| Error::MissingParameter("form body".to_string(), false))?;})
                .to_tokens(stream);
        }

        #[cfg(feature = "validate")]
        self.gen_validate_block(stream, &id, optional);
    }

    #[allow(dead_code)]
    #[allow(unused_imports)]
    fn gen_json_param(&self, stream: &mut TokenStream, optional: bool) {
        let id = Ident::new(self.name.as_str(), Span::call_site());
        let typ = self.typ.as_ref().expect("This should not happens");

        (quote! {
            // ================= \\\\\\\\
            let #id = req.parse_json::<#typ>().await
        })
        .to_tokens(stream);

        if optional {
            (quote! {.ok();}).to_tokens(stream);
        } else {
            (quote! {?;}).to_tokens(stream);
        }

        #[cfg(feature = "validate")]
        self.gen_validate_block(stream, &id, optional);
    }

    // =================
    #[cfg(feature = "validate")]
    fn gen_validate_block(&self, stream: &mut TokenStream, id: &Ident, optional: bool) {
        if self.validated {
            if optional {
                if self.is_vec {
                    (quote! {
                    if let Some(param) = &#id {
                        use ::validator::Validate;
                        for t in param.iter() {
                            t.validate().map_err(|e| Error::ValidationErrors(e))?;
                        }
                    }})
                    .to_tokens(stream);
                } else {
                    (quote! {
                    if let Some(param) = &#id {
                        use ::validator::Validate;
                        param.validate().map_err(|e| Error::ValidationErrors(e))?;
                    }})
                    .to_tokens(stream);
                }
            } else {
                if self.is_vec {
                    (quote! {
                    {   // =========
                        use ::validator::Validate;
                        for t in #id.iter() {
                            t.validate().map_err(|e| Error::ValidationErrors(e))?;
                        }
                    }})
                    .to_tokens(stream);
                } else {
                    (quote! {
                    {   // =============
                        use ::validator::Validate;
                        #id.validate().map_err(|e| Error::ValidationErrors(e))?;
                    }})
                    .to_tokens(stream);
                }
            }
        }
    }

    fn gen_cookie_param(&self, stream: &mut TokenStream) {
        let id = Ident::new(self.name.as_str(), Span::call_site());
        (quote! {

            let #id = req.take_cookies();
        })
        .to_tokens(stream);
    }

    fn gen_path_param(&self, stream: &mut TokenStream, optional: bool) {
        let name = self.name.as_str();
        let id = Ident::new(self.name.as_str(), Span::call_site());
        (quote! {

            let #id = req.params_mut().remove(#name)
        })
        .to_tokens(stream);

        self.gen_param_str_parsing(stream, name, optional);

        (quote! {;}).to_tokens(stream);
    }

    #[allow(dead_code)]
    #[allow(unused_imports)]
    fn gen_query(&self, stream: &mut TokenStream, optional: bool) {
        let id = Ident::new(self.name.as_str(), Span::call_site());
        let typ = self.typ.as_ref().expect("This should not happens");

        (quote! {
            let #id = req.parse_query::<#typ>().await
        })
        .to_tokens(stream);

        if optional {
            (quote! {.ok();}).to_tokens(stream);
        } else {
            (quote! {?;}).to_tokens(stream);
        }

        #[cfg(feature = "validate")]
        self.gen_validate_block(stream, &id, optional);
    }

    fn gen_query_param(&self, stream: &mut TokenStream, optional: bool) {
        let name = self.name.as_str();
        let id = Ident::new(self.name.as_str(), Span::call_site());
        (quote! {

            let #id = query.remove(#name)
        })
        .to_tokens(stream);

        self.gen_param_str_parsing(stream, name, optional);

        (quote! {;}).to_tokens(stream);
    }

    fn gen_param_str_parsing(&self, stream: &mut TokenStream, name: &str, optional: bool) {
        if !self.is_string() && self.typ.is_some() {
            let typ = self.typ.as_ref().unwrap(); // ==============================
            (quote! {.map(|p| p.parse::<#typ>()).transpose().map_err(|_| Error::InvalidParameter(#name.to_string(), false))?}).to_tokens(stream);
        }
        // =================================
        if !optional {
            (quote! {.ok_or_else(|| Error::MissingParameter(#name.to_string(), false))?})
                .to_tokens(stream);
        }
    }
}