use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, Arg, ItemFn};

pub fn hook_impl(attr: TokenStream, item: TokenStream) -> TokenStream {
    let args = parse_macro_input!(attr as HookArgs);
    let func = parse_macro_input!(item as ItemFn);

    let func_name = func.sig.ident.clone();
    let func_args = func.sig.inputs.clone();
    let func_ret = func.sig.output.clone();

    // Generate register extraction code
    let arg_extracts: Vec<_> = func_args
        .iter()
        .enumerate()
        .map(|(i, arg)| {
            let arg_name = match arg {
                Arg::Receiver(_) => return quote! { self },
                Arg::Typed(t) => &t.pat,
            };

            let reg = match i {
                0 => quote! { reg.rdi },
                1 => quote! { reg.rsi },
                2 => quote! { reg.rdx },
                3 => quote! { reg.rcx },
                4 => quote! { reg.r8 },
                5 => quote! { reg.r9 },
                _ => {
                    // Stack arguments
                    quote! { unsafe { reg.get_stack(#i - 6) } }
                }
            };

            quote! { #arg_name }
        })
        .collect();

    // Generate the hook handler
    let result = quote! {
        #[ctor::ctor]
        fn #func_name() {
            use prehook::hook::{HOOK_REGISTRY, HookInstance, HookVariant};
            use prehook::hook::RetnRoutine;

            let handler: RetnRoutine = {
                unsafe extern "win64" fn __hook_handler(
                    reg: *mut prehook::hook::Registers,
                    orig_func_ptr: usize,
                    _stack: usize,
                ) -> usize {
                    // SAFETY: We're in the hook handler, reg is valid
                    let reg = &mut *reg;

                    // Call the user's function with extracted arguments
                    let result = #func_name(#(#arg_extracts),*);

                    // For Retn hook, we return our result
                    // If user wants to call original, they use prehook::call_original!
                    result as usize
                }
                __hook_handler
            };

            let hook = HookInstance::new_retn(
                #args.offset,
                #args.symbol,
                stringify!(#func_name).to_string(),
                handler,
            );

            HOOK_REGISTRY.register(hook);
        }
    };

    result.into()
}

#[derive(syn::Parse)]
struct HookArgs {
    #[syn(attr)]
    offset: Option<syn::LitInt>,
    #[syn(attr)]
    symbol: Option<syn::LitStr>,
}

impl Parse for HookArgs {
    fn parse(input: parse::ParseStream) -> syn::Result<Self> {
        let mut offset = None;
        let mut symbol = None;

        while !input.is_empty() {
            let ident: syn::Ident = input.parse()?;
            let _: syn::Token![=] = input.parse()?;

            if ident == "offset" {
                offset = Some(input.parse()?);
            } else if ident == "symbol" {
                symbol = Some(input.parse()?);
            }
            if !input.is_empty() {
                let _: syn::Token![,] = input.parse()?;
            }
        }

        Ok(HookArgs { offset, symbol })
    }
}

mod parse {
    pub use syn::parse::Parse;
}
