subscription.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::ext::IdentExt;
use syn::{
    Block, Error, FnArg, ImplItem, ItemImpl, Pat, ReturnType, Type, TypeImplTrait, TypeParamBound,
    TypeReference,
};

use crate::args::{self, ComplexityType, RenameRuleExt, RenameTarget, SubscriptionField};
use crate::output_type::OutputType;
use crate::utils::{
    gen_deprecation, generate_default, generate_guards, get_cfg_attrs, get_crate_name, get_rustdoc,
    get_type_path_and_name, parse_complexity_expr, parse_graphql_attrs, remove_graphql_attrs,
    visible_fn, GeneratorResult,
};

pub fn generate(
    subscription_args: &args::Subscription,
    item_impl: &mut ItemImpl,
) -> GeneratorResult<TokenStream> {
    let crate_name = get_crate_name(subscription_args.internal);
    let (self_ty, self_name) = get_type_path_and_name(item_impl.self_ty.as_ref())?;
    let generics = &item_impl.generics;
    let where_clause = &item_impl.generics.where_clause;
    let extends = subscription_args.extends;

    let gql_typename = subscription_args
        .name
        .clone()
        .unwrap_or_else(|| RenameTarget::Type.rename(self_name.clone()));

    let desc = if subscription_args.use_type_description {
        quote! { ::std::option::Option::Some(::std::borrow::ToOwned::to_owned(<Self as #crate_name::Description>::description())) }
    } else {
        get_rustdoc(&item_impl.attrs)
            .map(|s| quote!(::std::option::Option::Some(::std::borrow::ToOwned::to_owned(#s))))
            .unwrap_or_else(|| quote!(::std::option::Option::None))
    };

    let mut create_stream = Vec::new();
    let mut schema_fields = Vec::new();

    for item in &mut item_impl.items {
        if let ImplItem::Fn(method) = item {
            let field: SubscriptionField = parse_graphql_attrs(&method.attrs)?.unwrap_or_default();
            if field.skip {
                remove_graphql_attrs(&mut method.attrs);
                continue;
            }

            let ident = &method.sig.ident;
            let field_name = field.name.clone().unwrap_or_else(|| {
                subscription_args
                    .rename_fields
                    .rename(method.sig.ident.unraw().to_string(), RenameTarget::Field)
            });
            let field_desc = get_rustdoc(&method.attrs)
                .map(|s| quote! {::std::option::Option::Some(::std::borrow::ToOwned::to_owned(#s))})
                .unwrap_or_else(|| quote! {::std::option::Option::None});
            let field_deprecation = gen_deprecation(&field.deprecation, &crate_name);
            let cfg_attrs = get_cfg_attrs(&method.attrs);

            if method.sig.asyncness.is_none() {
                return Err(Error::new_spanned(
                    &method,
                    "The subscription stream function must be asynchronous",
                )
                .into());
            }

            let ty = match &method.sig.output {
                ReturnType::Type(_, ty) => OutputType::parse(ty)?,
                ReturnType::Default => {
                    return Err(Error::new_spanned(
                        &method.sig.output,
                        "Resolver must have a return type",
                    )
                    .into())
                }
            };

            let mut create_ctx = true;
            let mut args = Vec::new();

            for (idx, arg) in method.sig.inputs.iter_mut().enumerate() {
                if let FnArg::Receiver(receiver) = arg {
                    if idx != 0 {
                        return Err(Error::new_spanned(
                            receiver,
                            "The self receiver must be the first parameter.",
                        )
                        .into());
                    }
                } else if let FnArg::Typed(pat) = arg {
                    if idx == 0 {
                        return Err(Error::new_spanned(
                            pat,
                            "The self receiver must be the first parameter.",
                        )
                        .into());
                    }

                    match (&*pat.pat, &*pat.ty) {
                        (Pat::Ident(arg_ident), Type::Path(arg_ty)) => {
                            args.push((
                                arg_ident.clone(),
                                arg_ty.clone(),
                                parse_graphql_attrs::<args::SubscriptionFieldArgument>(&pat.attrs)?
                                    .unwrap_or_default(),
                            ));
                            remove_graphql_attrs(&mut pat.attrs);
                        }
                        (arg, Type::Reference(TypeReference { elem, .. })) => {
                            if let Type::Path(path) = elem.as_ref() {
                                if idx != 1 || path.path.segments.last().unwrap().ident != "Context"
                                {
                                    return Err(Error::new_spanned(
                                        arg,
                                        "Only types that implement `InputType` can be used as input arguments.",
                                    )
                                    .into());
                                }
                                create_ctx = false;
                            }
                        }
                        _ => {
                            return Err(Error::new_spanned(arg, "Incorrect argument type").into());
                        }
                    }
                } else {
                    return Err(Error::new_spanned(arg, "Incorrect argument type").into());
                }
            }

            if create_ctx {
                let arg = syn::parse2::<FnArg>(quote! { _: &#crate_name::Context<'_> }).unwrap();
                method.sig.inputs.insert(1, arg);
            }

            let mut schema_args = Vec::new();
            let mut use_params = Vec::new();
            let mut get_params = Vec::new();

            for (
                ident,
                ty,
                args::SubscriptionFieldArgument {
                    name,
                    desc,
                    default,
                    default_with,
                    validator,
                    visible: arg_visible,
                    secret,
                },
            ) in &args
            {
                let name = name.clone().unwrap_or_else(|| {
                    subscription_args
                        .rename_args
                        .rename(ident.ident.unraw().to_string(), RenameTarget::Argument)
                });
                let desc = desc
                    .as_ref()
                    .map(|s| quote! {::std::option::Option::Some(::std::borrow::ToOwned::to_owned(#s))})
                    .unwrap_or_else(|| quote! {::std::option::Option::None});
                let default = generate_default(default, default_with)?;

                let schema_default = default
                    .as_ref()
                    .map(|value| {
                        quote! {
                            ::std::option::Option::Some(<#ty as #crate_name::LegacyInputType>::to_value(&#value))
                        }
                    })
                    .unwrap_or_else(|| quote! {::std::option::Option::None});

                let visible = visible_fn(arg_visible);
                schema_args.push(quote! {
                    args.insert(::std::borrow::ToOwned::to_owned(#name), #crate_name::registry::MetaInputValue {
                        name: ::std::borrow::ToOwned::to_owned(#name),
                        description: #desc,
                        ty: <#ty as #crate_name::LegacyInputType>::create_type_info(registry),
                        default_value: #schema_default,
                        validators: None,
                        visible: #visible,
                        is_secret: #secret,
                        rename: None
                    });
                });

                use_params.push(quote! { #ident });

                let default = match default {
                    Some(default) => quote! { ::std::option::Option::Some(|| -> #ty { #default }) },
                    None => quote! { ::std::option::Option::None },
                };
                let validators = validator.clone().unwrap_or_default().create_validators(
                    &crate_name,
                    quote!(&#ident),
                    quote!(#ty),
                    Some(quote!(.map_err(|err| err.into_server_error(__pos)))),
                )?;

                get_params.push(quote! {
                    #[allow(non_snake_case)]
                    let (__pos, #ident) = ctx.param_value::<#ty>(#name, #default)?;
                    #validators
                });
            }

            let res_ty = ty.value_type();
            let stream_ty = if let Type::ImplTrait(TypeImplTrait { bounds, .. }) = &res_ty {
                let mut r = None;
                for b in bounds {
                    if let TypeParamBound::Trait(b) = b {
                        r = Some(quote! { dyn #b });
                    }
                }
                quote! { #r }
            } else {
                quote! { #res_ty }
            };

            if let OutputType::Value(inner_ty) = &ty {
                let block = &method.block;
                let new_block = quote!({
                    {
                        let value = (move || { async move #block })().await;
                        ::std::result::Result::Ok(value)
                    }
                });
                method.block = syn::parse2::<Block>(new_block).expect("invalid block");
                method.sig.output =
                    syn::parse2::<ReturnType>(quote! { -> #crate_name::Result<#inner_ty> })
                        .expect("invalid result type");
            }

            let visible = visible_fn(&field.visible);
            let complexity = if let Some(complexity) = &field.complexity {
                match complexity {
                    ComplexityType::Const(n) => {
                        quote! { ::std::option::Option::Some(#crate_name::registry::ComplexityType::Const(#n)) }
                    }
                    ComplexityType::Fn(s) => {
                        let (variables, expr) = parse_complexity_expr(s)?;
                        let mut parse_args = Vec::new();
                        for variable in variables {
                            if let Some((
                                ident,
                                ty,
                                args::SubscriptionFieldArgument {
                                    name,
                                    default,
                                    default_with,
                                    ..
                                },
                            )) = args
                                .iter()
                                .find(|(pat_ident, _, _)| pat_ident.ident == variable)
                            {
                                let default = match generate_default(default, default_with)? {
                                    Some(default) => {
                                        quote! { ::std::option::Option::Some(|| -> #ty { #default }) }
                                    }
                                    None => quote! { ::std::option::Option::None },
                                };
                                let name = name.clone().unwrap_or_else(|| {
                                    subscription_args.rename_args.rename(
                                        ident.ident.unraw().to_string(),
                                        RenameTarget::Argument,
                                    )
                                });
                                parse_args.push(quote! {
                                    let #ident: #ty = __ctx.param_value(__variables_definition, __field, #name, #default)?;
                                });
                            }
                        }
                        quote! {
                            Some(#crate_name::registry::ComplexityType::Fn(|__ctx, __variables_definition, __field, child_complexity| {
                                #(#parse_args)*
                                ::std::result::Result::Ok(#expr)
                            }))
                        }
                    }
                }
            } else {
                quote! { ::std::option::Option::None }
            };

            schema_fields.push(quote! {
                #(#cfg_attrs)*
                fields.insert(::std::borrow::ToOwned::to_owned(#field_name), #crate_name::registry::MetaField {
                    name: ::std::borrow::ToOwned::to_owned(#field_name),
                    mapped_name: None,
                    description: #field_desc,
                    args: {
                        let mut args = #crate_name::indexmap::IndexMap::new();
                        #(#schema_args)*
                        args
                    },
                    ty: <<#stream_ty as #crate_name::futures_util::stream::Stream>::Item as #crate_name::LegacyOutputType>::create_type_info(registry),
                    deprecation: #field_deprecation,
                    cache_control: ::std::default::Default::default(),
                    external: false,
                    requires: ::std::option::Option::None,
                    provides: ::std::option::Option::None,
                    visible: #visible,
                    edges: Vec::new(),
                    relation: None,
                    compute_complexity: #complexity,
                    resolver: #crate_name::registry::resolvers::Resolver::Parent,
                    required_operation: None,
                    auth: None,
                });
            });

            let create_field_stream = quote! {
                self.#ident(ctx, #(#use_params),*)
                    .await
                    .map_err(|err| {
                        ::std::convert::Into::<#crate_name::Error>::into(err).into_server_error(ctx.item.pos)
                            .with_path(::std::vec![#crate_name::PathSegment::Field(::std::borrow::ToOwned::to_owned(&*field_name))])
                    })
            };

            let guard_map_err = quote! {
                .map_err(|err| {
                    err.into_server_error(ctx.item.pos)
                        .with_path(::std::vec![#crate_name::PathSegment::Field(::std::borrow::ToOwned::to_owned(&*field_name))])
                })
            };
            let guard = match &field.guard {
                Some(code) => Some(generate_guards(&crate_name, code, guard_map_err)?),
                None => None,
            };
            let stream_fn = quote! {
                let field_name = ::std::clone::Clone::clone(&ctx.item.node.response_key().node);
                let field = ::std::sync::Arc::new(::std::clone::Clone::clone(&ctx.item));

                let f = async {
                    #(#get_params)*
                    #guard
                    #create_field_stream
                };
                let stream = f.await.map_err(|err| ctx.set_error_path(err))?;

                let pos = ctx.item.pos;
                let schema_env = ::std::clone::Clone::clone(&ctx.schema_env);
                let query_env = ::std::clone::Clone::clone(&ctx.query_env);
                let stream = #crate_name::futures_util::stream::StreamExt::then(stream, {
                    let field_name = ::std::clone::Clone::clone(&field_name);
                    move |msg| {
                        let schema_env = ::std::clone::Clone::clone(&schema_env);
                        let query_env = ::std::clone::Clone::clone(&query_env);
                        let field = ::std::clone::Clone::clone(&field);
                        let field_name = ::std::clone::Clone::clone(&field_name);
                        async move {
                            let ctx_selection_set = query_env.create_context(
                                &schema_env,
                                ::std::option::Option::Some(#crate_name::QueryPathNode {
                                    parent: ::std::option::Option::None,
                                    segment: #crate_name::QueryPathSegment::Name(&field_name),
                                }),
                                None,
                                &field.node.selection_set,
                            );

                            let type_name = <<#stream_ty as #crate_name::futures_util::stream::Stream>::Item as #crate_name::LegacyOutputType>::qualified_type_name();
                            let mut execute_fut = async {
                                #[allow(bare_trait_objects)]
                                let ri = #crate_name::extensions::ResolveInfo {
                                    path_node: ctx_selection_set.path_node.as_ref().unwrap(),
                                    parent_type: #gql_typename,
                                    return_type: &<<#stream_ty as #crate_name::futures_util::stream::Stream>::Item as #crate_name::LegacyOutputType>::qualified_type_name(),
                                    name: field.node.name.node.as_str(),
                                    alias: field.node.alias.as_ref().map(|alias| alias.node.as_str()),
                                    required_operation: None,
                                    auth: None,
                                    input_values: Vec::new() // Not needed
                                };

                                let resolve_fut = async {
                                    let resolved = #crate_name::LegacyOutputType::resolve(&msg, &ctx_selection_set, &*field).await?;
                                    Ok(Some(resolved))
                                };

                                /*
                                let resolve_fut = async {
                                    #crate_name::LegacyOutputType::resolve(&msg, &ctx_selection_set, &*field)
                                        .await
                                        .map(::std::option::Option::Some)
                                };
                                */

                                #crate_name::futures_util::pin_mut!(resolve_fut);

                                let mut resp = match query_env.extensions.resolve(ri, &mut resolve_fut).await {
                                    Ok(Some(value)) => {
                                        let mut container = #crate_name::graph_entities::ResponseContainer::new_container();
                                        container.insert(#crate_name::graph_entities::ResponseNodeRelation::NotARelation { response_key: None, field: ::std::clone::Clone::clone(&field_name).to_string().into() }, value);
                                        let response = &mut *ctx.response_graph.write().await;
                                        let container_id = response.insert_node(container);
                                        response.set_root_unchecked(container_id);
                                        let a = #crate_name::graph_entities::QueryResponse::default();
                                        let b = std::mem::replace(response, a);
                                        #crate_name::Response::new(b, #crate_name::parser::types::OperationType::Subscription)
                                    }
                                    Ok(None) => {
                                        let mut container = #crate_name::graph_entities::ResponseContainer::new_container();
                                        let response = &mut *ctx.response_graph.write().await;
                                        let primitive_node_id = response.insert_node(Box::new(#crate_name::graph_entities::ResponsePrimitive::default()));
                                        container.insert(#crate_name::graph_entities::ResponseNodeRelation::NotARelation { response_key: None, field: ::std::clone::Clone::clone(&field_name).to_string().into() }, primitive_node_id);
                                        let container_id = response.insert_node(container);
                                        response.set_root_unchecked(container_id);
                                        let a = #crate_name::graph_entities::QueryResponse::default();
                                        let b = std::mem::replace(response, a);
                                        #crate_name::Response::new(b, #crate_name::parser::types::OperationType::Subscription)
                                    }
                                    Err(err) => #crate_name::Response::from_errors(vec![err], #crate_name::parser::types::OperationType::Subscription),
                                };

                                use ::std::iter::Extend;
                                resp.errors.extend(::std::mem::take(&mut *query_env.errors.lock().unwrap()));
                                resp
                            };
                            #crate_name::futures_util::pin_mut!(execute_fut);
                            ::std::result::Result::Ok(query_env.extensions.execute(query_env.operation_name.as_deref(), &query_env.operation, &mut execute_fut).await)
                        }
                    }
                });
                #crate_name::ServerResult::Ok(stream)
            };

            create_stream.push(quote! {
                #(#cfg_attrs)*
                if ctx.item.node.name.node == #field_name {
                    let stream = #crate_name::futures_util::stream::TryStreamExt::try_flatten(
                        #crate_name::futures_util::stream::once((move || async move { #stream_fn })())
                    );
                    let stream = #crate_name::futures_util::StreamExt::map(stream, |res| match res {
                        ::std::result::Result::Ok(resp) => resp,
                        ::std::result::Result::Err(err) => #crate_name::Response::from_errors(::std::vec![err], #crate_name::parser::types::OperationType::Subscription),
                    });
                    return ::std::option::Option::Some(::std::boxed::Box::pin(stream));
                }
            });

            remove_graphql_attrs(&mut method.attrs);
        }
    }

    if create_stream.is_empty() {
        return Err(Error::new_spanned(
            self_ty,
            "A GraphQL Object type must define one or more fields.",
        )
        .into());
    }

    let visible = visible_fn(&subscription_args.visible);

    let expanded = quote! {
        #item_impl

        #[allow(clippy::all, clippy::pedantic)]
        #[allow(unused_braces, unused_variables)]
        impl #generics #crate_name::SubscriptionType for #self_ty #where_clause {
            fn type_name() -> ::std::borrow::Cow<'static, ::std::primitive::str> {
                ::std::borrow::Cow::Borrowed(#gql_typename)
            }

            #[allow(bare_trait_objects)]
            fn create_type_info(registry: &mut #crate_name::registry::Registry) -> #crate_name::registry::InputValueType {
                registry.create_subscription_type::<Self, _>(|registry|
                    #crate_name::registry::MetaType::Object(#crate_name::registry::ObjectType {
                        name: ::std::borrow::ToOwned::to_owned(#gql_typename),
                        description: #desc,
                        fields: {
                            let mut fields = #crate_name::indexmap::IndexMap::new();
                            #(#schema_fields)*
                            fields
                        },
                        cache_control: ::std::default::Default::default(),
                        extends: #extends,
                        keys: ::std::option::Option::None,
                        visible: #visible,
                        is_subscription: true,
                        is_node: false,
                        rust_typename: ::std::borrow::ToOwned::to_owned(::std::any::type_name::<Self>()),
                        constraints: vec![],
                    })
                ).into()
            }

            fn create_field_stream<'__life>(
                &'__life self,
                ctx: &'__life #crate_name::Context<'_>,
            ) -> ::std::option::Option<::std::pin::Pin<::std::boxed::Box<dyn #crate_name::futures_util::stream::Stream<Item = #crate_name::Response> + ::std::marker::Send + '__life>>> {
                #(#create_stream)*
                ::std::option::Option::None
            }
        }
    };

    Ok(expanded.into())
}