extern crate proc_macro;
use proc_macro::TokenStream;
use quote::quote;
use syn::{
    Attribute, Block, Error, GenericParam, Generics, Ident, ItemStruct, ItemTrait, Signature,
    Token, Type, Visibility,
    parse::{Parse, ParseStream, Result},
    parse_macro_input, parse_quote,
    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<Rsc, WL, Tag> for WL {
            #(#impls)*
        }
    }
    .into()
}

#[proc_macro_derive(DefaultUiState, attributes(default_ui_state))]
pub fn derive_default_ui_state(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 {
        if !found_attr
            && let Type::Path(path) = &field.ty
            && path.path.is_ident("DefaultUiState")
        {
            state_field = Some(field);
        }
        let Some(attr) = field
            .attrs
            .iter()
            .find(|a| a.path().is_ident("default_ui_state"))
        else {
            continue;
        };
        if found_attr {
            output.extend(
                Error::new(
                    attr.span(),
                    "cannot have more than one default_ui_state 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(), "no DefaultUiState field found").into_compile_error(),
        );
        return output.into();
    };
    let sname = &state.ident;
    let fname = field.ident.as_ref().unwrap();
    output.extend(quote! {
        impl iris::default::HasDefaultUiState for #sname {
            fn default_state(&self) -> &iris::default::DefaultUiState {
                &self.#fname
            }
            fn default_state_mut(&mut self) -> &mut iris::default::DefaultUiState {
                &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()
}