// SPDX-License-Identifier: Apache-2.0

//! Procedural macro crate for `substrait-validator-core`.
//!
//! The derive macros defined here are essentially an ugly workaround for the
//! lack of any protobuf introspection functionality provided by prost.
//! Basically, they take (the AST of) the code generated by prost and try to
//! recover the needed protobuf message metadata from there. Things would have
//! been a *LOT* simpler and a *LOT* less brittle if prost would simply
//! provide this information via traits of its own, but alas, there doesn't
//! seem to be a way to do this without forking prost, and introspection
//! seems to be a non-goal of that project.
//!
//! Besides being ugly, this method is rather brittle and imprecise when it
//! comes to recovering field names, due to the various case conversions
//! automatically done by protoc and prost. Some known issues are:
//!
//!  - The recovered type name for messages defined within messages uses
//!    incorrect case conventions, as the procedural macros have no way
//!    of distinguishing packages from message definition scopes in the
//!    type path.
//!  - If the .proto source files use unexpected case conventions for
//!    various things, the resulting case conventions for types, field names,
//!    oneof variants, and enum variants will be wrong.
//!  - Whenever the .proto source files name a field using something that is
//!    a reserved word in Rust (notably `type`), prost will use a raw
//!    identifier to represent the name. This syntax is currently not filtered
//!    out from the recovered names, so a field named `type` becomes `r#type`.
//!    This is probably not a fundamental problem, though.
//!
//! Ultimately, however, these names are only used for diagnostic messages and
//! the likes. In the worst case, the above inconsistencies may confuse the
//! user, but they should not affect the valid/invalid/maybe-valid result of
//! the validator or cause compile- or runtime errors.

extern crate proc_macro;

use heck::{ToShoutySnakeCase, ToSnakeCase};
use proc_macro::TokenStream;
use quote::quote;

/// Converts a Rust identifier string generated via stringify!() to the
/// original identifier by "cooking" raw identifiers.
fn cook_ident(ident: &syn::Ident) -> String {
    let ident = ident.to_string();
    if let Some((_, keyword)) = ident.split_once('#') {
        keyword.to_string()
    } else {
        ident
    }
}

#[doc(hidden)]
#[proc_macro_derive(ProtoMeta, attributes(proto_meta))]
pub fn proto_meta(input: TokenStream) -> TokenStream {
    proto_meta_derive(syn::parse_macro_input!(input))
}

fn proto_meta_derive(ast: syn::DeriveInput) -> TokenStream {
    match ast.data {
        syn::Data::Struct(ref struct_data) => proto_meta_derive_message(&ast, struct_data),
        syn::Data::Enum(ref enum_data) => match enum_data.variants.iter().next().unwrap().fields {
            syn::Fields::Unit => {
                for variant in enum_data.variants.iter() {
                    if !matches!(variant.fields, syn::Fields::Unit) {
                        panic!("all variants of a protobuf oneof enum must have a single, unnamed field");
                    }
                }

                proto_meta_derive_enum(&ast, enum_data)
            }
            syn::Fields::Unnamed(..) => {
                for variant in enum_data.variants.iter() {
                    if let syn::Fields::Unnamed(fields) = &variant.fields {
                        if fields.unnamed.len() != 1 {
                            panic!("all variants of a protobuf oneof enum must have a single, unnamed field");
                        }
                    } else {
                        panic!("all variants of a protobuf oneof enum must have a single, unnamed field");
                    }
                }

                proto_meta_derive_oneof(&ast, enum_data)
            }
            _ => panic!("enum with named elements don't map to protobuf constructs"),
        },
        syn::Data::Union(_) => panic!("unions don't map to protobuf constructs"),
    }
}

enum FieldType {
    Optional,
    BoxedOptional,
    Repeated,
    Primitive,
}

fn is_repeated(typ: &syn::Type) -> FieldType {
    if let syn::Type::Path(path) = typ {
        if let Some(last) = path.path.segments.last() {
            if last.ident == "Option" {
                if let syn::PathArguments::AngleBracketed(ref args) = last.arguments {
                    if let syn::GenericArgument::Type(syn::Type::Path(path2)) =
                        args.args.first().unwrap()
                    {
                        if path2.path.segments.last().unwrap().ident == "Box" {
                            return FieldType::BoxedOptional;
                        } else {
                            return FieldType::Optional;
                        }
                    }
                }
                panic!("Option without type argument?");
            } else if last.ident == "Vec" {
                if let syn::PathArguments::AngleBracketed(ref args) = last.arguments {
                    if let syn::GenericArgument::Type(syn::Type::Path(path2)) =
                        args.args.first().unwrap()
                    {
                        if path2.path.segments.last().unwrap().ident == "u8" {
                            return FieldType::Primitive;
                        } else {
                            return FieldType::Repeated;
                        }
                    }
                }
                panic!("Vec without type argument?");
            }
        }
    }
    FieldType::Primitive
}

fn proto_meta_derive_message(ast: &syn::DeriveInput, data: &syn::DataStruct) -> TokenStream {
    let name = &ast.ident;
    let name_str = cook_ident(name);
    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();

    let parse_unknown_matches: Vec<_> = data
        .fields
        .iter()
        .map(|field| {
            if let Some(ident) = &field.ident {
                let ident_str = cook_ident(ident);
                let action = match is_repeated(&field.ty) {
                    FieldType::Optional => quote! {
                        crate::parse::traversal::push_proto_field(
                            y,
                            &self.#ident.as_ref(),
                            #ident_str,
                            true,
                            |_, _| Ok(()),
                        );
                    },
                    FieldType::BoxedOptional => quote! {
                        crate::parse::traversal::push_proto_field(
                            y,
                            &self.#ident,
                            #ident_str,
                            true,
                            |_, _| Ok(()),
                        );
                    },
                    FieldType::Repeated => quote! {
                        crate::parse::traversal::push_proto_repeated_field(
                            y,
                            &self.#ident.as_ref(),
                            #ident_str,
                            true,
                            |_, _| Ok(()),
                            |_, _, _, _, _| (),
                        );
                    },
                    FieldType::Primitive => quote! {
                        use crate::input::traits::ProtoPrimitive;
                        if !y.config.ignore_unknown_fields || !self.#ident.proto_primitive_is_default() {
                            crate::parse::traversal::push_proto_field(
                                y,
                                &Some(&self.#ident),
                                #ident_str,
                                true,
                                |_, _| Ok(()),
                            );
                        }
                    },
                };
                quote! {
                    if !y.field_parsed(#ident_str) {
                        unknowns = true;
                        #action
                    }
                }
            } else {
                panic!("protobuf message fields must have names");
            }
        })
        .collect();

    quote!(
        impl #impl_generics crate::input::traits::ProtoMessage for #name #ty_generics #where_clause {
            fn proto_message_type() -> &'static str {
                use ::once_cell::sync::Lazy;
                static TYPE_NAME: Lazy<::std::string::String> = Lazy::new(|| {
                    crate::input::proto::cook_path(module_path!(), #name_str)
                });
                &TYPE_NAME
            }
        }

        impl #impl_generics crate::input::traits::InputNode for #name #ty_generics #where_clause {
            fn type_to_node() -> crate::output::tree::Node {
                use crate::input::traits::ProtoMessage;
                crate::output::tree::NodeType::ProtoMessage(Self::proto_message_type()).into()
            }

            fn data_to_node(&self) -> crate::output::tree::Node {
                use crate::input::traits::ProtoMessage;
                crate::output::tree::NodeType::ProtoMessage(Self::proto_message_type()).into()
            }

            fn oneof_variant(&self) -> Option<&'static str> {
                None
            }

            fn parse_unknown(
                &self,
                y: &mut crate::parse::context::Context<'_>,
            ) -> bool {
                let mut unknowns = false;
                #(#parse_unknown_matches)*
                unknowns
            }
        }
    )
    .into()
}

fn proto_meta_derive_oneof(ast: &syn::DeriveInput, data: &syn::DataEnum) -> TokenStream {
    let name = &ast.ident;
    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();

    let variant_matches: Vec<_> = data
        .variants
        .iter()
        .map(|variant| {
            let ident = &variant.ident;
            let proto_name = cook_ident(ident).to_snake_case();
            quote! { #name::#ident (_) => #proto_name }
        })
        .collect();

    let node_matches: Vec<_> = data
        .variants
        .iter()
        .map(|variant| {
            let ident = &variant.ident;
            quote! { #name::#ident (x) => x.data_to_node() }
        })
        .collect();

    let parse_unknown_matches: Vec<_> = data
        .variants
        .iter()
        .map(|variant| {
            let ident = &variant.ident;
            quote! { #name::#ident (x) => x.parse_unknown(y) }
        })
        .collect();

    quote!(
        impl #impl_generics crate::input::traits::ProtoOneOf for #name #ty_generics #where_clause {
            fn proto_oneof_variant(&self) -> &'static str {
                match self {
                    #(#variant_matches),*
                }
            }
        }

        impl #impl_generics crate::input::traits::InputNode for #name #ty_generics #where_clause {
            fn type_to_node() -> crate::output::tree::Node {
                crate::output::tree::NodeType::ProtoMissingOneOf.into()
            }

            fn data_to_node(&self) -> crate::output::tree::Node {
                match self {
                    #(#node_matches),*
                }
            }

            fn oneof_variant(&self) -> Option<&'static str> {
                use crate::input::traits::ProtoOneOf;
                Some(self.proto_oneof_variant())
            }

            fn parse_unknown(
                &self,
                y: &mut crate::parse::context::Context<'_>,
            ) -> bool {
                match self {
                    #(#parse_unknown_matches),*
                }
            }
        }
    )
    .into()
}

fn proto_meta_derive_enum(ast: &syn::DeriveInput, data: &syn::DataEnum) -> TokenStream {
    let name = &ast.ident;
    let name_str = cook_ident(name);
    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();

    let upper_name = name_str.to_shouty_snake_case();

    let variant_names: Vec<_> = data
        .variants
        .iter()
        .map(|variant| {
            let ident = &variant.ident;
            let proto_name = format!(
                "{}_{}",
                upper_name,
                cook_ident(ident).to_shouty_snake_case()
            );
            (ident, proto_name)
        })
        .collect();

    let variant_matches: Vec<_> = variant_names
        .iter()
        .map(|(ident, proto_name)| {
            quote! { #name::#ident => #proto_name }
        })
        .collect();

    let (_, first_variant_name) = &variant_names[0];

    quote!(
        impl #impl_generics crate::input::traits::ProtoEnum for #name #ty_generics #where_clause {
            fn proto_enum_type() -> &'static str {
                use ::once_cell::sync::Lazy;
                static TYPE_NAME: Lazy<::std::string::String> = Lazy::new(|| {
                    crate::input::proto::cook_path(module_path!(), #name_str)
                });
                &TYPE_NAME
            }

            fn proto_enum_default_variant() -> &'static str {
                #first_variant_name
            }

            fn proto_enum_variant(&self) -> &'static str {
                match self {
                    #(#variant_matches),*
                }
            }

            fn proto_enum_from_i32(x: i32) -> Option<Self> {
                Self::from_i32(x)
            }
        }
    )
    .into()
}