zus-macros 1.1.4

use {
  proc_macro::TokenStream,
  quote::quote,
  syn::{
    Expr, ExprLit, ItemImpl, ItemStruct, Lit, Meta, Token, Type,
    parse::{Parse, ParseStream},
    parse_macro_input,
    punctuated::Punctuated,
  },
};

/// Configuration options for zus_service
#[derive(Default)]
struct ServiceConfig {
  name: Option<String>,
  port: Option<u16>,
  host: Option<String>,
  compression: Option<String>,
  max_connections: Option<usize>,
}

impl Parse for ServiceConfig {
  fn parse(input: ParseStream) -> syn::Result<Self> {
    let mut config = ServiceConfig::default();

    if input.is_empty() {
      return Ok(config);
    }

    let args = Punctuated::<Meta, Token![,]>::parse_terminated(input)?;

    for arg in args {
      if let Meta::NameValue(nv) = arg {
        let ident = nv
          .path
          .get_ident()
          .ok_or_else(|| syn::Error::new_spanned(&nv.path, "Expected identifier"))?;

        match ident.to_string().as_str() {
          | "name" => {
            if let Expr::Lit(ExprLit {
              lit: Lit::Str(lit_str), ..
            }) = &nv.value
            {
              config.name = Some(lit_str.value());
            }
          }
          | "port" => {
            if let Expr::Lit(ExprLit {
              lit: Lit::Int(lit_int), ..
            }) = &nv.value
            {
              config.port = Some(lit_int.base10_parse()?);
            }
          }
          | "host" => {
            if let Expr::Lit(ExprLit {
              lit: Lit::Str(lit_str), ..
            }) = &nv.value
            {
              config.host = Some(lit_str.value());
            }
          }
          | "compression" => {
            if let Expr::Lit(ExprLit {
              lit: Lit::Str(lit_str), ..
            }) = &nv.value
            {
              config.compression = Some(lit_str.value());
            }
          }
          | "max_connections" => {
            if let Expr::Lit(ExprLit {
              lit: Lit::Int(lit_int), ..
            }) = &nv.value
            {
              config.max_connections = Some(lit_int.base10_parse()?);
            }
          }
          | _ => {
            return Err(syn::Error::new_spanned(
              &nv.path,
              format!("Unknown configuration option: {ident}"),
            ));
          }
        }
      }
    }

    Ok(config)
  }
}

/// Attribute macro to mark a struct as a ZUS service
///
/// This automatically implements the Service trait and handles method routing.
///
/// # Example
/// ```ignore
/// #[zus_service]
/// impl MyService {
///     #[method("SayHello")]
///     async fn say_hello(&self, request: HelloRequest) -> Result<HelloResponse> {
///         // Implementation
///     }
/// }
///
/// // With configuration:
/// #[zus_service(name = "MyService", port = 9527, host = "0.0.0.0")]
/// impl MyService {
///     #[method("SayHello")]
///     async fn say_hello(&self, request: HelloRequest) -> Result<HelloResponse> {
///         // Implementation
///     }
/// }
/// ```
#[proc_macro_attribute]
pub fn zus_service(args: TokenStream, input: TokenStream) -> TokenStream {
  // Parse configuration
  let config = parse_macro_input!(args as ServiceConfig);
  let input_impl = parse_macro_input!(input as ItemImpl);

  // Extract the struct name
  let struct_name = match &*input_impl.self_ty {
    | Type::Path(type_path) => &type_path.path.segments.last().unwrap().ident,
    | _ => {
      return syn::Error::new_spanned(&input_impl.self_ty, "Expected a simple type path")
        .to_compile_error()
        .into();
    }
  };

  // Collect all methods with #[method] attribute
  let mut method_arms = Vec::new();
  let mut impl_items = Vec::new();

  for item in input_impl.items {
    if let syn::ImplItem::Fn(mut method) = item {
      // Check if it has #[method] attribute
      let mut method_name = None;
      method.attrs.retain(|attr| {
        if attr.path().is_ident("method") {
          if let Ok(Lit::Str(lit_str)) = attr.parse_args::<Lit>() {
            method_name = Some(lit_str.value());
          }
          false // Remove the attribute
        } else {
          true // Keep other attributes
        }
      });

      if let Some(rpc_method_name) = method_name {
        let fn_name = &method.sig.ident;

        // Generate match arm for this method
        method_arms.push(quote! {
            #rpc_method_name => {
                self.#fn_name(params, context).await
            }
        });
      }

      impl_items.push(syn::ImplItem::Fn(method));
    } else {
      impl_items.push(item);
    }
  }

  // Use configured service name or default to struct name
  let service_name = config.name.unwrap_or_else(|| struct_name.to_string());

  // Reconstruct the impl block with processed items
  let self_ty = &input_impl.self_ty;
  let generics = &input_impl.generics;
  let where_clause = &input_impl.generics.where_clause;

  // Generate config constants
  let config_impl = if config.port.is_some()
    || config.host.is_some()
    || config.compression.is_some()
    || config.max_connections.is_some()
  {
    let port = config.port.unwrap_or(9527);
    let host = config.host.unwrap_or_else(|| "0.0.0.0".to_string());
    let compression = config.compression.unwrap_or_else(|| "quicklz".to_string());
    let max_connections = config.max_connections.unwrap_or(1000);

    quote! {
        impl #struct_name {
            /// Default port for this service (from macro configuration)
            pub const DEFAULT_PORT: u16 = #port;

            /// Default host for this service (from macro configuration)
            pub const DEFAULT_HOST: &'static str = #host;

            /// Default compression algorithm (from macro configuration)
            pub const DEFAULT_COMPRESSION: &'static str = #compression;

            /// Default max connections (from macro configuration)
            pub const DEFAULT_MAX_CONNECTIONS: usize = #max_connections;

            /// Create a pre-configured server manager for this service
            pub fn create_server() -> zus_rpc_server::ZusServerManager {
                zus_rpc_server::ZusServerManager::new(
                    Self::DEFAULT_HOST.to_string(),
                    Self::DEFAULT_PORT
                )
                .with_max_connections(Self::DEFAULT_MAX_CONNECTIONS)
            }
        }
    }
  } else {
    quote! {}
  };

  // Generate the Service implementation
  let expanded = quote! {
      // Original impl block with processed methods
      impl #generics #self_ty #where_clause {
          #(#impl_items)*
      }

      // Configuration constants and helper methods
      #config_impl

      #[async_trait::async_trait]
      impl zus_rpc_server::Service for #struct_name {
          fn service_name(&self) -> &str {
              #service_name
          }

          async fn do_work(
              &self,
              method: &str,
              params: bytes::Bytes,
              context: zus_rpc_server::RequestContext,
          ) -> zus_common::Result<bytes::Bytes> {
              match method {
                  #(#method_arms)*
                  _ => Err(zus_common::ZusError::MethodNotFound(
                      format!("Unknown method: {}", method)
                  ))
              }
          }
      }
  };

  TokenStream::from(expanded)
}

/// Attribute macro to mark a method as a ZUS RPC method
///
/// This is used in conjunction with #[zus_service].
///
/// # Example
/// ```ignore
/// #[method("SayHello")]
/// async fn say_hello(&self, request: HelloRequest) -> Result<HelloResponse> { }
/// ```
#[proc_macro_attribute]
pub fn method(_args: TokenStream, input: TokenStream) -> TokenStream {
  // This is a marker attribute that is processed by #[zus_service]
  // Just return the input as-is
  input
}

/// Derive macro for ZUS services
///
/// Automatically implements common service functionality.
///
/// # Example
/// ```ignore
/// #[derive(ZusService)]
/// struct MyService {
///     config: ServiceConfig,
/// }
/// ```
#[proc_macro_derive(ZusService, attributes(service_name))]
pub fn derive_zus_service(input: TokenStream) -> TokenStream {
  let input = parse_macro_input!(input as ItemStruct);
  let name = &input.ident;

  // Extract service_name attribute if present
  let service_name = input
    .attrs
    .iter()
    .find(|attr| attr.path().is_ident("service_name"))
    .and_then(|attr| {
      if let Meta::NameValue(ref meta) = attr.meta
        && let Expr::Lit(ExprLit {
          lit: Lit::Str(ref lit_str),
          ..
        }) = meta.value
      {
        return Some(lit_str.value());
      }
      None
    })
    .unwrap_or_else(|| name.to_string());

  let expanded = quote! {
      impl #name {
          pub const SERVICE_NAME: &'static str = #service_name;
          pub const IS_ZUS_SERVICE: bool = true;
      }
  };

  TokenStream::from(expanded)
}

/// Macro to create a ZUS server with default configuration
///
/// This simplifies server setup by providing sensible defaults.
///
/// # Example
/// ```ignore
/// zus_server! {
///     host: "0.0.0.0",
///     port: 9527,
///     services: [MyService::new(), AnotherService::new()],
/// }
/// ```
#[proc_macro]
pub fn zus_server(_input: TokenStream) -> TokenStream {
  // For now, provide a simple expansion
  // A full implementation would parse the configuration
  let expanded = quote! {
      compile_error!("zus_server! macro not yet fully implemented. Use ZusServerManager directly.");
  };

  TokenStream::from(expanded)
}