use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, Attribute, Data, DeriveInput, Fields, Ident};

pub(crate) fn rich_error_impl(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let enum_ident = input.ident; // InnerError
    let vis = input.vis; // e.g. pub

    let outer_ident = match parse_outer_ident(&input.attrs) {
        Ok(Some(ident)) => ident,
        Ok(None) => format_ident!("Error"),
        Err(err) => return err.to_compile_error().into(),
    };

    let data_enum = match input.data {
        Data::Enum(e) => e,
        _ => {
            return syn::Error::new_spanned(enum_ident, "RichError can only be derived for enums")
                .to_compile_error()
                .into()
        }
    };

    let mut ctor_fns = Vec::new();
    let mut from_wrapped_impls = Vec::new();

    for variant in data_enum.variants {
        let var_ident = variant.ident;
        let fn_ident = to_snake_case_ident(&var_ident);

        match variant.fields {
            // Struct variant: Variant { field: Ty, ... }
            Fields::Named(fields_named) => {
                let field_idents: Vec<_> = fields_named
                    .named
                    .iter()
                    .map(|f| f.ident.clone().unwrap())
                    .collect();

                let field_types: Vec<_> = fields_named.named.iter().map(|f| f.ty.clone()).collect();

                let ctor = quote! {
                    #[track_caller]
                    #vis fn #fn_ident( #( #field_idents : #field_types ),* ) -> #outer_ident {
                        #outer_ident::new(
                            #enum_ident::#var_ident { #( #field_idents ),* }
                        )
                    }
                };
                ctor_fns.push(ctor);

                // If the struct variant has exactly one field, also generate From<FieldTy> for Error
                if field_types.len() == 1 {
                    let ty = &field_types[0];
                    let arg = &field_idents[0];
                    let impl_block = quote! {
                        impl ::core::convert::From<#ty> for #outer_ident {
                            #[track_caller]
                            fn from(#arg: #ty) -> Self {
                                #outer_ident::#fn_ident(#arg)
                            }
                        }
                    };
                    from_wrapped_impls.push(impl_block);
                }
            }

            // Tuple variant: Variant(Ty0, Ty1, ...)
            Fields::Unnamed(fields_unnamed) => {
                let field_count = fields_unnamed.unnamed.len();

                let field_idents: Vec<_> = (0..field_count)
                    .map(|i| format_ident!("arg{}", i))
                    .collect();

                let field_types: Vec<_> = fields_unnamed
                    .unnamed
                    .iter()
                    .map(|f| f.ty.clone())
                    .collect();

                let ctor = quote! {
                    #[track_caller]
                    #vis fn #fn_ident( #( #field_idents : #field_types ),* ) -> #outer_ident {
                        #outer_ident::new(
                            #enum_ident::#var_ident( #( #field_idents ),* )
                        )
                    }
                };
                ctor_fns.push(ctor);

                // If there is exactly one field, generate From<FieldTy> for Error
                if field_types.len() == 1 {
                    let ty = &field_types[0];
                    let arg = &field_idents[0];
                    let impl_block = quote! {
                        impl ::core::convert::From<#ty> for #outer_ident {
                            #[track_caller]
                            fn from(#arg: #ty) -> Self {
                                #outer_ident::#fn_ident(#arg)
                            }
                        }
                    };
                    from_wrapped_impls.push(impl_block);
                }
            }

            // Unit variant: Variant
            Fields::Unit => {
                let ctor = quote! {
                    #[track_caller]
                    #vis fn #fn_ident() -> #outer_ident {
                        #outer_ident::new(#enum_ident::#var_ident)
                    }
                };
                ctor_fns.push(ctor);
            }
        }
    }

    let expanded = quote! {
        // Outer error type: wraps InnerError and stores creation location
        #[derive(Debug)]
        #vis struct #outer_ident {
            inner: #enum_ident,
            location: &'static ::std::panic::Location<'static>,
        }

        // Display prints the inner error plus the source location
        impl ::std::fmt::Display for #outer_ident {
            fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
                write!(
                    f,
                    "{} (at {}:{})",
                    self.inner,
                    self.location.file(),
                    self.location.line(),
                )
            }
        }

        // Implement std::error::Error, delegating source() to InnerError
        impl ::std::error::Error for #outer_ident {
            fn source(&self) -> ::std::option::Option<&(dyn ::std::error::Error + 'static)> {
                ::std::option::Option::Some(&self.inner)
            }
        }

        impl #outer_ident {
            #[track_caller]
            #vis fn new(inner: #enum_ident) -> Self {
                Self {
                    inner,
                    location: ::std::panic::Location::caller(),
                }
            }

            #( #ctor_fns )*

            #vis fn inner(&self) -> &#enum_ident {
                &self.inner
            }

            #vis fn location(&self) -> &'static ::std::panic::Location<'static> {
                self.location
            }
        }

        // From<InnerError> for Error so `?` on Result<_, InnerError> works
        impl ::core::convert::From<#enum_ident> for #outer_ident {
            #[track_caller]
            fn from(inner: #enum_ident) -> Self {
                #outer_ident::new(inner)
            }
        }

        // From<WrappedError> for Error for single-field variants
        #( #from_wrapped_impls )*
    };

    TokenStream::from(expanded)
}

fn to_snake_case_ident(ident: &syn::Ident) -> syn::Ident {
    let s = ident.to_string();
    let mut out = String::new();
    for (i, ch) in s.chars().enumerate() {
        if ch.is_uppercase() {
            if i != 0 {
                out.push('_');
            }
            out.extend(ch.to_lowercase());
        } else {
            out.push(ch);
        }
    }
    format_ident!("{}", out)
}

fn parse_outer_ident(attrs: &[Attribute]) -> syn::Result<Option<Ident>> {
    let mut outer_ident = None;

    for attr in attrs {
        if !attr.path().is_ident("rich_error") {
            continue;
        }

        attr.parse_nested_meta(|meta| {
            if meta.path.is_ident("outer") {
                if outer_ident.is_some() {
                    return Err(meta.error("`outer` specified more than once"));
                }
                let value: Ident = meta.value()?.parse()?;
                outer_ident = Some(value);
                Ok(())
            } else {
                Err(meta.error("unsupported rich_error option"))
            }
        })?;
    }

    Ok(outer_ident)
}
