extern crate proc_macro;
use proc_macro::TokenStream;
use quote::quote;
use syn::{
    Attribute, Block, Error, FnArg, GenericParam, Generics, Ident, ItemFn, ItemStruct, ItemTrait,
    ReturnType, Signature, Token, Type, Visibility,
    parse::{Parse, ParseStream, Result},
    parse_macro_input, parse_quote,
    spanned::Spanned,
};

/// Marks the initializer called when Android creates an Iris view.
///
/// An attribute is necessary here because the Android loader requires one
/// exported `JNI_OnLoad` symbol and `android-view` requires a plain function
/// pointer monomorphized for the application state. A function returning a
/// custom state remains its factory; a function with no return value receives
/// `&mut AndroidUiState` and uses that state directly. The generated linker and
/// JNI glue is Android-gated; the annotated function therefore does not need
/// its own `cfg` attribute.
#[proc_macro_attribute]
pub fn android_init(args: TokenStream, item: TokenStream) -> TokenStream {
    if !args.is_empty() {
        return Error::new(
            proc_macro2::Span::call_site(),
            "android_init takes no arguments",
        )
        .into_compile_error()
        .into();
    }

    let function = parse_macro_input!(item as ItemFn);
    let name = &function.sig.ident;
    let (state, direct_initializer): (Type, bool) = match &function.sig.output {
        ReturnType::Default => (parse_quote!(::iris::android::AndroidUiState), true),
        ReturnType::Type(_, state) => ((**state).clone(), false),
    };
    if function.sig.inputs.len() != 2
        || function
            .sig
            .inputs
            .iter()
            .any(|argument| !matches!(argument, FnArg::Typed(_)))
    {
        return Error::new(
            function.sig.inputs.span(),
            "an android_init function takes UI state and resources",
        )
        .into_compile_error()
        .into();
    }
    if function.sig.asyncness.is_some()
        || function.sig.constness.is_some()
        || matches!(function.sig.safety, syn::Safety::Unsafe(_))
        || !function.sig.generics.params.is_empty()
    {
        return Error::new(
            function.sig.span(),
            "an android_init function must be a plain, non-generic synchronous function",
        )
        .into_compile_error()
        .into();
    }

    let factory = if direct_initializer {
        quote! {
            fn init(
                mut ui_state: ::iris::android::AndroidUiState,
                rsc: &mut <#state as ::iris::android::AndroidAppState>::Resources,
            ) -> #state {
                super::#name(&mut ui_state, rsc);
                ui_state
            }
        }
    } else {
        quote! {}
    };
    let create = if direct_initializer {
        quote! { init }
    } else {
        quote! { super::#name }
    };

    quote! {
        #[cfg(target_os = "android")]
        #function

        #[cfg(target_os = "android")]
        mod __iris_android_app {
            use super::*;

            #factory

            extern "system" fn new_view_peer<'local>(
                env: ::iris::android::__private::JNIEnv<'local>,
                view: ::iris::android::__private::View<'local>,
                context: ::iris::android::__private::Context<'local>,
            ) -> ::iris::android::__private::JLong {
                ::iris::android::new_peer::<#state>(env, view, context, #create)
            }

            #[unsafe(no_mangle)]
            pub unsafe extern "system" fn JNI_OnLoad(
                vm: *mut ::iris::android::__private::RawJavaVM,
                _: *mut ::core::ffi::c_void,
            ) -> ::iris::android::__private::JInt {
                unsafe { ::iris::android::__private::on_load(vm, new_view_peer) }
            }
        }
    }
    .into()
}

struct Input {
    attrs: Vec<Attribute>,
    vis: Visibility,
    name: Ident,
    generics: Generics,
    fns: Vec<InputFn>,
}

struct InputFn {
    attrs: Vec<Attribute>,
    sig: Signature,
    body: Block,
}

impl Parse for Input {
    fn parse(input: ParseStream) -> Result<Self> {
        let attrs = input.call(Attribute::parse_outer)?;
        let vis = input.parse()?;
        input.parse::<Token![trait]>()?;
        let name = input.parse()?;
        let generics = input.parse::<Generics>()?;
        input.parse::<Token![;]>()?;
        let mut fns = Vec::new();
        while !input.is_empty() {
            let attrs = input.call(Attribute::parse_outer)?;
            let sig = input.parse()?;
            let body = input.parse()?;
            fns.push(InputFn { attrs, sig, body })
        }
        if !input.is_empty() {
            input.error("function expected");
        }
        Ok(Input {
            attrs,
            vis,
            name,
            generics,
            fns,
        })
    }
}

#[proc_macro]
pub fn widget_trait(input: TokenStream) -> TokenStream {
    let Input {
        attrs,
        vis,
        name,
        mut generics,
        fns,
    } = parse_macro_input!(input as Input);

    let sigs: Vec<_> = fns
        .iter()
        .map(|InputFn { attrs, sig, .. }| quote! { #(#attrs)* #sig })
        .collect();
    let impls: Vec<_> = fns
        .iter()
        .map(|InputFn { sig, body, .. }| quote! { #sig #body })
        .collect();

    let Some(GenericParam::Type(state)) = generics.params.first() else {
        return Error::new(name.span(), "expected state generic parameter")
            .into_compile_error()
            .into();
    };

    let state = &state.ident;

    generics
        .params
        .push(parse_quote!(WL: WidgetLike<#state, Tag>));
    generics.params.push(parse_quote!(Tag));

    let mut trai: ItemTrait = parse_quote!(
        #vis trait #name #generics {
            #(#sigs;)*
        }
    );

    trai.attrs = attrs;

    quote! {
        #trai

        impl #generics #name<Rsc, WL, Tag> for WL {
            #(#impls)*
        }
    }
    .into()
}

#[proc_macro_derive(DesktopUiState, attributes(desktop_ui_state))]
pub fn derive_desktop_ui_state(input: TokenStream) -> TokenStream {
    let state: ItemStruct = parse_macro_input!(input);
    derive_ui_state(
        state,
        UiStateDerive {
            module: "desktop",
            state_type: "DesktopUiState",
            state_trait: "HasDesktopUiState",
            field_attr: "desktop_ui_state",
            get: "desktop_state",
            get_mut: "desktop_state_mut",
        },
    )
}

#[proc_macro_derive(AndroidUiState, attributes(android_ui_state))]
pub fn derive_android_ui_state(input: TokenStream) -> TokenStream {
    let state: ItemStruct = parse_macro_input!(input);
    derive_ui_state(
        state,
        UiStateDerive {
            module: "android",
            state_type: "AndroidUiState",
            state_trait: "HasAndroidUiState",
            field_attr: "android_ui_state",
            get: "android_state",
            get_mut: "android_state_mut",
        },
    )
}

struct UiStateDerive {
    module: &'static str,
    state_type: &'static str,
    state_trait: &'static str,
    field_attr: &'static str,
    get: &'static str,
    get_mut: &'static str,
}

fn derive_ui_state(state: ItemStruct, names: UiStateDerive) -> TokenStream {
    let UiStateDerive {
        module,
        state_type,
        state_trait,
        field_attr,
        get,
        get_mut,
    } = names;
    let mut output = proc_macro2::TokenStream::new();
    let mut found_attr = false;
    let mut state_field = None;
    for field in &state.fields {
        if !found_attr
            && let Type::Path(path) = &field.ty
            && path.path.is_ident(state_type)
        {
            state_field = Some(field);
        }
        let Some(attr) = field.attrs.iter().find(|a| a.path().is_ident(field_attr)) else {
            continue;
        };
        if found_attr {
            output.extend(
                Error::new(
                    attr.span(),
                    format!("cannot have more than one {field_attr} attribute"),
                )
                .into_compile_error(),
            );
            continue;
        }
        found_attr = true;
        state_field = Some(field);
    }
    let Some(field) = state_field else {
        output.extend(
            Error::new(state.ident.span(), format!("no {state_type} field found"))
                .into_compile_error(),
        );
        return output.into();
    };
    let sname = &state.ident;
    let Some(fname) = field.ident.as_ref() else {
        return Error::new(
            field.span(),
            format!("the {state_type} field must be named"),
        )
        .into_compile_error()
        .into();
    };
    let module = Ident::new(module, sname.span());
    let state_type = Ident::new(state_type, sname.span());
    let state_trait = Ident::new(state_trait, sname.span());
    let get = Ident::new(get, sname.span());
    let get_mut = Ident::new(get_mut, sname.span());
    let (impl_generics, type_generics, where_clause) = state.generics.split_for_impl();
    output.extend(quote! {
        impl #impl_generics iris::#module::#state_trait for #sname #type_generics #where_clause {
            fn #get(&self) -> &iris::#module::#state_type {
                &self.#fname
            }
            fn #get_mut(&mut self) -> &mut iris::#module::#state_type {
                &mut self.#fname
            }
        }
    });
    output.into()
}

#[proc_macro_derive(WidgetView, attributes(root))]
pub fn derive_widget_view(input: TokenStream) -> TokenStream {
    let mut output = proc_macro2::TokenStream::new();

    let state: ItemStruct = parse_macro_input!(input);

    let mut found_attr = false;
    let mut state_field = None;
    for field in &state.fields {
        let Some(attr) = field.attrs.iter().find(|a| a.path().is_ident("root")) else {
            continue;
        };
        if found_attr {
            output.extend(
                Error::new(attr.span(), "cannot have more than one root widget")
                    .into_compile_error(),
            );
            continue;
        }
        found_attr = true;
        state_field = Some(field);
    }
    let Some(field) = state_field else {
        output.extend(
            Error::new(state.ident.span(), "no root widget field found (#[root])")
                .into_compile_error(),
        );
        return output.into();
    };
    let sname = &state.ident;
    let fname = field.ident.as_ref().unwrap();
    let fty = &field.ty;
    output.extend(quote! {
        impl iris::core::WidgetView for #sname {
            type Widget = <#fty as iris::core::HasWidget>::Widget;
            fn root(&self) -> #fty {
                self.#fname
            }
        }
    });
    output.into()
}