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

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

struct InputFn {
    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 sig = input.parse()?;
            let body = input.parse()?;
            fns.push(InputFn { 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(|f| f.sig.clone()).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<State, WL, Tag> for WL {
            #(#impls)*
        }
    }
    .into()
}

#[proc_macro_derive(GlobalState, attributes(has))]
pub fn derive_global_state(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemStruct);
    let name = input.ident;
    let mut impls = TokenStream::new();
    for field in input.fields {
        let Some(attr) = field.attrs.iter().find(|a| a.path().is_ident("has")) else {
            continue;
        };
        let error: TokenStream = Error::new(
            attr.span(),
            "invalid attribute format; usage: #[has(HasTrait, trait_fn)]",
        )
        .into_compile_error()
        .into();
        let Meta::List(list) = &attr.meta else {
            return error;
        };
        match list.parse_args_with(Punctuated::<Ident, Token![,]>::parse_terminated) {
            Ok(list) => {
                if list.len() != 2 {
                    return error;
                }
                let traitt = &list[0];
                let fn_name = &list[1];
                let field_name = field.ident;
                let ty = field.ty;
                impls.extend::<TokenStream>(
                    quote! {
                        impl #traitt for #name {
                            fn #fn_name(&mut self) -> &mut #ty {
                                &mut self.#field_name
                            }
                        }
                    }
                    .into(),
                );
            }
            Err(..) => {
                return error;
            }
        }
    }
    impls
}

#[proc_macro_derive(HasUi)]
pub fn derive_has_ui(input: TokenStream) -> TokenStream {
    has_ui(&parse_macro_input!(input))
}

fn has_ui(input: &ItemStruct) -> TokenStream {
    let name = &input.ident;
    let Some(field) = input
        .fields
        .iter()
        .find(|f| f.ty == parse_quote!(Ui<Self>) || f.ty == parse_quote!(Ui))
    else {
        return Error::new(name.span(), "could not find a Ui<Self> field for HasUi")
            .into_compile_error()
            .into();
    };
    let field = &field.ident;
    quote! {
        impl HasUi for #name {
            fn ui_ref(&self) -> &iris::iris_core::Ui<Self> {
                &self.#field
            }

            fn ui(&mut self) -> &mut iris::iris_core::Ui<Self> {
                &mut self.#field
            }
        }
    }
    .into()
}

#[proc_macro_derive(DefaultUiState)]
pub fn derive_default_ui_state(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemStruct);
    let mut output = has_ui(&input);
    let name = input.ident;
    let Some(field) = input
        .fields
        .iter()
        .find(|f| f.ty == parse_quote!(UiState<Self>) || f.ty == parse_quote!(UiState))
    else {
        return Error::new(
            name.span(),
            "could not find a UiState<Self> field for HasUiState",
        )
        .into_compile_error()
        .into();
    };
    let field = &field.ident;
    output.extend::<TokenStream>(
        quote! {
            impl HasUiState for #name {
                fn ui_state(&mut self) -> &mut iris::default::UiState<Self> {
                    &mut self.#field
                }
            }
        }
        .into(),
    );
    output
}