// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! PDK Unit Test Macros
//!
//! This module provides macros to simplify implementing GrpcBackend for unit tests.

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{parse_macro_input, FnArg, ImplItem, ItemImpl, Lit, Meta, NestedMeta, Pat, Type};

/// Attribute macro for implementing `GrpcBackend` on an impl block.
///
/// Annotate methods with `#[grpc_method(service = "ServiceName", method = "MethodName")]`
/// to automatically route gRPC calls to the appropriate handler.
///
/// # Example
///
/// ```rust,ignore
/// #[derive(Default)]
/// pub struct BackService {}
///
/// #[grpc_backend]
/// impl BackService {
///     #[grpc_method(service = "ExampleService", method = "ExampleMethod")]
///     fn example_method(&self, req: ExampleRequest) -> ExampleResponse {
///         ExampleResponse {
///             response: "body".to_string(),
///             ..Default::default()
///         }
///     }
/// }
/// ```
#[proc_macro_attribute]
pub fn protobuf_grpc_backend(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemImpl);

    match generate_grpc_backend_impl(input) {
        Ok(tokens) => TokenStream::from(tokens),
        Err(err) => TokenStream::from(err.to_compile_error()),
    }
}

struct GrpcMethodInfo {
    method_name: syn::Ident,
    service: String,
    grpc_method: String,
    request_type: Type,
}

fn generate_grpc_backend_impl(input: ItemImpl) -> syn::Result<TokenStream2> {
    let struct_type = &input.self_ty;
    let mut methods: Vec<GrpcMethodInfo> = Vec::new();

    // Collect all methods annotated with #[grpc_method]
    for item in &input.items {
        if let ImplItem::Method(method) = item {
            let mut service = None;
            let mut grpc_method = None;

            for attr in &method.attrs {
                if attr.path.is_ident("grpc_method") {
                    let meta = attr.parse_meta()?;
                    if let Meta::List(list) = meta {
                        for nested in list.nested {
                            if let NestedMeta::Meta(Meta::NameValue(nv)) = nested {
                                if nv.path.is_ident("service") {
                                    if let Lit::Str(lit) = nv.lit {
                                        service = Some(lit.value());
                                    }
                                } else if nv.path.is_ident("method") {
                                    if let Lit::Str(lit) = nv.lit {
                                        grpc_method = Some(lit.value());
                                    }
                                }
                            }
                        }
                    }
                }
            }

            if let (Some(svc), Some(mth)) = (service, grpc_method) {
                // Extract request type from the method signature
                let request_type = extract_request_type(method)?;

                methods.push(GrpcMethodInfo {
                    method_name: method.sig.ident.clone(),
                    service: svc,
                    grpc_method: mth,
                    request_type,
                });
            }
        }
    }

    // Generate match arms for each method
    let match_arms: Vec<TokenStream2> = methods.iter().map(|info| {
        let method_name = &info.method_name;
        let service = &info.service;
        let grpc_method = &info.grpc_method;
        let request_type = &info.request_type;
        quote! {
            (#service, #grpc_method) => {
                let Ok(request) = <#request_type as ::protobuf::Message>::parse_from_reader(&mut cursor) else {
                    return ::pdk_unit::UnitGrpcResponse::default().with_status_code(3).with_status("Deserialization failed.")
                };
                let response = self.#method_name(request);
                let Ok(bytes) = ::protobuf::Message::write_to_bytes(&response) else {
                    return ::pdk_unit::UnitGrpcResponse::default().with_status_code(3).with_status("Serialization failed.");
                };
                bytes
            }
        }
    }).collect();

    // Strip #[grpc_method] attributes from the original impl
    let mut cleaned_impl = input.clone();
    for item in &mut cleaned_impl.items {
        if let ImplItem::Method(method) = item {
            method
                .attrs
                .retain(|attr| !attr.path.is_ident("grpc_method"));
        }
    }

    let expanded = quote! {
        #cleaned_impl

        impl ::pdk_unit::GrpcBackend for #struct_type {
            fn call(&self, req: ::pdk_unit::UnitGrpcRequest) -> ::pdk_unit::UnitGrpcResponse {
                let data = req.message().cloned().unwrap_or_default();
                let mut cursor = ::std::io::Cursor::new(data);

                let bytes = match (req.service(), req.method()) {
                    #(#match_arms)*
                    (service, method) => {
                        return ::pdk_unit::UnitGrpcResponse::default()
                        .with_status_code(12)
                        .with_status(format!("Service: {} Method: {} not implemented", service, method))
                    }
                };

                ::pdk_unit::UnitGrpcResponse::default()
                    .with_status_code(0)
                    .with_message(bytes)
                    .with_status("OK")
            }
        }
    };

    Ok(expanded)
}

fn extract_request_type(method: &syn::ImplItemMethod) -> syn::Result<Type> {
    // Find the first non-self argument
    for arg in &method.sig.inputs {
        if let FnArg::Typed(pat_type) = arg {
            // Skip if it's self
            if let Pat::Ident(ident) = &*pat_type.pat {
                if ident.ident == "self" {
                    continue;
                }
            }
            return Ok((*pat_type.ty).clone());
        }
    }

    Err(syn::Error::new_spanned(
        &method.sig,
        "grpc_method must have a request parameter",
    ))
}
