use proc_macro::TokenStream;
use quote::{quote, format_ident};
use syn::{parse_macro_input, ItemTrait, TraitItem, FnArg, ReturnType, Type, PathArguments, GenericArgument};

pub fn derive_tunnel_service(attr: TokenStream, input: TokenStream) -> TokenStream {
    // Parse mandatory trait-level transaction type and optional DB binding:
    // #[tunnel_service(TxType)] or #[tunnel_service(TxType @ DB)]
    struct TunnelAttr {
        tx_type: Type,
        db_binding: Option<syn::Path>,
    }
    impl syn::parse::Parse for TunnelAttr {
        fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
            let tx_type: Type = input.parse()?;
            let db_binding = if input.peek(syn::Token![@]) {
                input.parse::<syn::Token![@]>()?;
                Some(input.parse::<syn::Path>()?)
            } else {
                None
            };
            Ok(TunnelAttr { tx_type, db_binding })
        }
    }
    let TunnelAttr { tx_type, db_binding } = syn::parse(attr)
        .expect("tunnel_service: must specify transaction type, e.g. #[tunnel_service(MyTx)] or #[tunnel_service(MyTx @ DB)]");

    let input_trait = parse_macro_input!(input as ItemTrait);
    let trait_name = &input_trait.ident;
    let trait_vis = &input_trait.vis;
    let service_name_str = trait_name.to_string();

    let proxy_name = format_ident!("{}Proxy", trait_name);
    let dispatcher_name = format_ident!("{}Dispatcher", trait_name);

    // Collect method info
    let mut methods = Vec::new();
    for item in &input_trait.items {
        if let TraitItem::Fn(method) = item {
            let method_name = &method.sig.ident;
            let method_name_str = method_name.to_string();
            let full_method_name = format!("{}.{}", service_name_str, method_name_str);

            // Extract return type: Result<T, RpcError<E>> -> (T, E)
            let (ok_type, err_type) = extract_result_types(&method.sig.output);

            // Extract parameters (skip &self)
            let mut param_names = Vec::new();
            let mut param_types = Vec::new();
            for arg in &method.sig.inputs {
                if let FnArg::Typed(pat_type) = arg {
                    param_names.push(&*pat_type.pat);
                    param_types.push(&*pat_type.ty);
                }
            }

            methods.push(MethodInfo {
                method_name: method_name.clone(),
                full_method_name,
                ok_type,
                err_type,
                param_names,
                param_types,
            });
        }
    }

    // Generate trait with #[async_trait] and async methods
    let trait_attrs = &input_trait.attrs;
    let trait_supertraits = &input_trait.supertraits;
    let trait_methods: Vec<_> = methods.iter().map(|m| {
        let name = &m.method_name;
        let ok_type = &m.ok_type;
        let err_type = &m.err_type;
        let params: Vec<_> = m.param_names.iter().zip(m.param_types.iter()).map(|(n, t)| {
            quote! { #n: #t }
        }).collect();

        // All methods get scope injected with trait-level tx type
        let tx_ty = &tx_type;
        quote! {
            async fn #name(&self, scope: std::sync::Arc<mad_hatter::TxScope<#tx_ty>>, #(#params),*) -> Result<#ok_type, mad_hatter::tunnel::RpcError<#err_type>>;
        }
    }).collect();

    let trait_output = quote! {
        #(#trait_attrs)*
        #[::async_trait::async_trait]
        #trait_vis trait #trait_name: #trait_supertraits {
            #(#trait_methods)*
        }
    };

    // Generate Proxy — always impl ProxyName directly (not impl Trait)
    // Proxy methods never have scope parameter (remote calls don't need transactions)
    let proxy_methods: Vec<_> = methods.iter().map(|m| {
        let name = &m.method_name;
        let full_name = &m.full_method_name;
        let ok_type = &m.ok_type;
        let err_type = &m.err_type;
        let param_names = &m.param_names;
        let param_types = &m.param_types;

        let params: Vec<_> = param_names.iter().zip(param_types.iter()).map(|(n, t)| {
            quote! { #n: #t }
        }).collect();

        let serialize_params = if param_names.is_empty() {
            quote! { serde_json::Value::Null }
        } else if param_names.len() == 1 {
            let p = &param_names[0];
            quote! {
                serde_json::to_value(&#p)
                    .map_err(|e| mad_hatter::tunnel::RpcError::<#err_type>::Transport(
                        format!("serialize params: {}", e)))?
            }
        } else {
            quote! {
                serde_json::to_value(&(#(&#param_names),*))
                    .map_err(|e| mad_hatter::tunnel::RpcError::<#err_type>::Transport(
                        format!("serialize params: {}", e)))?
            }
        };

        quote! {
            pub async fn #name(&self, #(#params),*) -> Result<#ok_type, mad_hatter::tunnel::RpcError<#err_type>> {
                let params = #serialize_params;
                match self.__caller.call(#full_name, params).await {
                    Ok(v) => {
                        serde_json::from_value(v)
                            .map_err(|e| mad_hatter::tunnel::RpcError::Transport(
                                format!("deserialize response: {}", e)))
                    }
                    Err(mad_hatter::tunnel::TunnelError::Service(v)) => {
                        let err: #err_type = serde_json::from_value(v)
                            .map_err(|e| mad_hatter::tunnel::RpcError::Transport(
                                format!("deserialize error: {}", e)))?;
                        Err(mad_hatter::tunnel::RpcError::Service(err))
                    }
                    Err(mad_hatter::tunnel::TunnelError::Transport(msg)) => {
                        Err(mad_hatter::tunnel::RpcError::Transport(msg))
                    }
                }
            }
        }
    }).collect();

    let proxy_output = quote! {
        #trait_vis struct #proxy_name {
            __caller: std::sync::Arc<dyn mad_hatter::tunnel::TunnelCaller>,
        }

        impl #proxy_name {
            pub fn new(caller: std::sync::Arc<dyn mad_hatter::tunnel::TunnelCaller>) -> Self {
                Self { __caller: caller }
            }

            #(#proxy_methods)*
        }
    };

    // Generate Dispatcher
    let dispatch_arms: Vec<_> = methods.iter().map(|m| {
        let name = &m.method_name;
        let full_name = &m.full_method_name;
        let param_names = &m.param_names;
        let param_types = &m.param_types;

        let deserialize_params = if param_names.is_empty() {
            quote! {}
        } else if param_names.len() == 1 {
            let p = &param_names[0];
            let t = &param_types[0];
            quote! {
                let #p: #t = serde_json::from_value(params.clone())
                    .map_err(|e| serde_json::to_value(e.to_string()).unwrap_or(serde_json::Value::Null))?;
            }
        } else {
            let indices: Vec<_> = (0..param_names.len()).map(syn::Index::from).collect();
            quote! {
                let __tuple: (#(#param_types),*) = serde_json::from_value(params.clone())
                    .map_err(|e| serde_json::to_value(e.to_string()).unwrap_or(serde_json::Value::Null))?;
                #(let #param_names = __tuple.#indices;)*
            }
        };

        // All methods wrapped in tunnel_tx_guard with trait-level tx type
        let tx_ty = &tx_type;
        let dispatch_arm = if let Some(ref db_path) = db_binding {
            quote! {
                #full_name => {
                    #deserialize_params
                    let handler = self.__handler.clone();
                    match mad_hatter::tunnel::tunnel_tx_guard_with_db::<#tx_ty, _>(
                        #db_path.get().expect("DB not initialized").db().clone(),
                        |scope| async move {
                            match handler.#name(scope, #(#param_names),*).await {
                                Ok(v) => serde_json::to_value(&v)
                                    .map_err(|e| serde_json::to_value(e.to_string()).unwrap_or(serde_json::Value::Null)),
                                Err(mad_hatter::tunnel::RpcError::Service(e)) => {
                                    Err(serde_json::to_value(&e).unwrap_or(serde_json::Value::Null))
                                }
                                Err(mad_hatter::tunnel::RpcError::Transport(msg)) => {
                                    Err(serde_json::to_value(&msg).unwrap_or(serde_json::Value::Null))
                                }
                            }
                        },
                    ).await {
                        Ok(inner) => inner,
                        Err(abort_value) => Err(abort_value),
                    }
                }
            }
        } else {
            quote! {
                #full_name => {
                    #deserialize_params
                    let handler = self.__handler.clone();
                    match mad_hatter::tunnel::tunnel_tx_guard::<#tx_ty, _>(|scope| async move {
                        match handler.#name(scope, #(#param_names),*).await {
                            Ok(v) => serde_json::to_value(&v)
                                .map_err(|e| serde_json::to_value(e.to_string()).unwrap_or(serde_json::Value::Null)),
                            Err(mad_hatter::tunnel::RpcError::Service(e)) => {
                                Err(serde_json::to_value(&e).unwrap_or(serde_json::Value::Null))
                            }
                            Err(mad_hatter::tunnel::RpcError::Transport(msg)) => {
                                Err(serde_json::to_value(&msg).unwrap_or(serde_json::Value::Null))
                            }
                        }
                    }).await {
                        Ok(inner) => inner,
                        Err(abort_value) => Err(abort_value),
                    }
                }
            }
        };
        dispatch_arm
    }).collect();

    let dispatcher_output = quote! {
        #trait_vis struct #dispatcher_name<T: #trait_name + 'static> {
            __handler: std::sync::Arc<T>,
        }

        impl<T: #trait_name + 'static> #dispatcher_name<T> {
            pub fn new(handler: std::sync::Arc<T>) -> Self {
                Self { __handler: handler }
            }

            pub fn boxed(handler: std::sync::Arc<T>) -> Box<dyn mad_hatter::tunnel::Dispatch> {
                Box::new(Self::new(handler))
            }
        }

        #[::async_trait::async_trait]
        impl<T: #trait_name + 'static> mad_hatter::tunnel::Dispatch for #dispatcher_name<T> {
            fn service_name(&self) -> &str {
                #service_name_str
            }

            async fn dispatch(&self, method: &str, params: serde_json::Value) -> Result<serde_json::Value, serde_json::Value> {
                match method {
                    #(#dispatch_arms)*
                    _ => Err(serde_json::to_value(format!("unknown method: {}", method))
                        .unwrap_or(serde_json::Value::Null)),
                }
            }
        }
    };

    let output = quote! {
        #trait_output
        #proxy_output
        #dispatcher_output
    };

    output.into()
}

struct MethodInfo<'a> {
    method_name: syn::Ident,
    full_method_name: String,
    ok_type: Box<Type>,
    err_type: Box<Type>,
    param_names: Vec<&'a syn::Pat>,
    param_types: Vec<&'a Type>,
}



/// Extract (T, E) from `Result<T, RpcError<E>>`.
///
/// Panics if the return type is not in the expected form.
fn extract_result_types(return_type: &ReturnType) -> (Box<Type>, Box<Type>) {
    if let ReturnType::Type(_, ty) = return_type {
        if let Type::Path(type_path) = ty.as_ref() {
            if let Some(segment) = type_path.path.segments.last() {
                if segment.ident == "Result" {
                    if let PathArguments::AngleBracketed(args) = &segment.arguments {
                        let mut args_iter = args.args.iter();

                        // First arg: T (ok type)
                        let ok_type = match args_iter.next() {
                            Some(GenericArgument::Type(t)) => Box::new(t.clone()),
                            _ => panic!("tunnel_service: Result must have Ok type"),
                        };

                        // Second arg: RpcError<E>
                        let err_type = match args_iter.next() {
                            Some(GenericArgument::Type(Type::Path(err_path))) => {
                                let err_segment = err_path.path.segments.last()
                                    .expect("tunnel_service: expected RpcError<E>");
                                if err_segment.ident != "RpcError" {
                                    panic!(
                                        "tunnel_service: methods must return Result<T, RpcError<E>>, \
                                         found Result<T, {}>",
                                        err_segment.ident
                                    );
                                }
                                if let PathArguments::AngleBracketed(rpc_args) = &err_segment.arguments {
                                    match rpc_args.args.first() {
                                        Some(GenericArgument::Type(e)) => Box::new(e.clone()),
                                        _ => panic!("tunnel_service: RpcError must have type parameter E"),
                                    }
                                } else {
                                    panic!("tunnel_service: RpcError must have type parameter <E>");
                                }
                            }
                            _ => panic!(
                                "tunnel_service: methods must return Result<T, RpcError<E>>"
                            ),
                        };

                        return (ok_type, err_type);
                    }
                }
            }
        }
    }
    panic!("tunnel_service: methods must return Result<T, RpcError<E>>");
}

