use std::collections::HashSet;

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::{
    parenthesized,
    parse::{Parse, ParseStream},
    parse_macro_input, Expr, FnArg, ItemFn, LitStr, Pat, Signature, Token,
};

#[proc_macro_attribute]
pub fn trace_method(attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parsing macros attributes
    let TraceAttrs {
        level,
        skip_self,
        skip_retval,
        skip,
    } = parse_macro_input!(attr as TraceAttrs);
    // Parsing passed method
    let func = parse_macro_input!(item as ItemFn);
    let vis = &func.vis;
    let sig = &func.sig;
    let name = &sig.ident;
    let attrs = &func.attrs;
    let body = &func.block;
    // Building called trace log fmt
    let (fmt_lit, fmt_args) = collect_called_log_parts(sig, skip_self, &skip);
    // Generating the wrapped function code
    let ret_log = if skip_retval {
        quote! {
            crate::logging::LOGGER_INSTANCE.write_log(
                #level, &format!("RET#{:04x} {}::{}(…)", __trace_id, __trace_typename, stringify!(#name))
            );
        }
    } else {
        quote! {
            crate::logging::LOGGER_INSTANCE.write_log(
                #level,
                &format!("RET#{:04x} {}::{}(…) → {:?}", __trace_id, __trace_typename, stringify!(#name), &__trace_ret)
            );
        }
    };
    let out = quote! {
        #(#attrs)*
        #vis #sig {
            if crate::logging::LOGGER_INSTANCE.should_log(#level) {
                let __trace_id = crate::get_trace_id!();
                let __trace_typename_full = ::std::any::type_name::<Self>();
                let __trace_typename = {
                    let s = __trace_typename_full;
                    // Find the first '<' to determine the end bound
                    let end = s.find('<').unwrap_or(s.len());
                    // Find the last '::' before '<'
                    let start = s[..end].rfind("::").map_or(0, |i| i + 2);
                    s[start..end].trim_end_matches('>')
                };
                crate::logging::LOGGER_INSTANCE.write_log(
                    #level, &format!(#fmt_lit, __trace_id, __trace_typename, stringify!(#name), #(#fmt_args),*)
                );
                let __trace_ret = (|| #body)();
                #ret_log;
                __trace_ret
            } else {
                (|| #body)()
            }
        }
    };
    return out.into();
}

fn collect_called_log_parts(
    sig: &Signature,
    skip_self: bool,
    skip: &HashSet<String>,
) -> (LitStr, Vec<proc_macro2::TokenStream>) {
    let mut fmt_called = String::with_capacity(sig.inputs.len() * 64 + 128);
    let mut fmt_args = Vec::with_capacity(sig.inputs.len());
    fmt_called.push_str("CALL#{:04x} {}::{}(");
    let mut added_arg_count = 0;
    for arg in &sig.inputs {
        match arg {
            FnArg::Receiver(_) => {
                if !skip_self {
                    fmt_called.push_str("self: {:?}, ");
                    fmt_args.push(quote! { self });
                    added_arg_count += 1;
                }
            }
            FnArg::Typed(pat_ty) => {
                if let Pat::Ident(id) = &*pat_ty.pat {
                    let ident = &id.ident;
                    let ident_str = ident.to_string();
                    if !skip.contains(&ident_str) {
                        fmt_called.push_str(&format!("{ident_str}: {{:?}}, "));
                        fmt_args.push(quote! { #ident });
                        added_arg_count += 1;
                    }
                }
            }
        }
    }
    if added_arg_count > 0 {
        fmt_called.pop();
        fmt_called.pop();
    }
    fmt_called.push(')');
    let fmt_lit = LitStr::new(&fmt_called, Span::call_site());
    return (fmt_lit, fmt_args);
}

struct TraceAttrs {
    level: Expr,
    skip_self: bool,
    skip_retval: bool,
    skip: HashSet<String>,
}

impl Parse for TraceAttrs {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let level: Expr = input.parse()?;
        let mut skip_self = false;
        let mut skip_retval = false;
        let mut skip = HashSet::new();
        if input.peek(Token![,]) {
            let _comma: Token![,] = input.parse()?;
            if !input.is_empty() {
                let skip_ident: syn::Ident = input.parse()?;
                if skip_ident != "skip" {
                    return Err(syn::Error::new_spanned(skip_ident, r#"expected `skip(...)`"#));
                }
                let content;
                parenthesized!(content in input);
                while !content.is_empty() {
                    if content.peek(Token![self]) {
                        let _kw_self: Token![self] = content.parse()?;
                        skip_self = true;
                    } else {
                        let id: syn::Ident = content.parse()?;
                        let s = id.to_string();
                        if s == "retval" {
                            skip_retval = true;
                        } else {
                            skip.insert(s);
                        }
                    }
                    if content.peek(Token![,]) {
                        let _ = content.parse::<Token![,]>()?;
                    } else {
                        break;
                    }
                }
            }
        }
        Ok(Self {
            level,
            skip_self,
            skip_retval,
            skip,
        })
    }
}
