mynt 0.1.0

use std::collections::HashMap;

use mynt::*;
use ordered_multimap::ListOrderedMultimap;
use proc_macro2::TokenStream;
use quote::{ToTokens, format_ident, quote, quote_spanned};

// declare the attribute macro
// (we could also wrap the macro with mynt!, but this allows for defining the macro in another file)
mynt_macro_attribute!(schema => schema_impl);

fn schema_impl(attr: TokenStream, input: TokenStream) -> TokenStream {
    mynt_assert!(attr.is_empty()); // ensure the outer macro doesn't have any arguments

    let syn::ItemMod {
        attrs,
        vis,
        unsafety,
        mod_token,
        ident,
        content,
        semi,
    } = syn::parse2(input).unwrap_or_quit(); // parse with syn or quit with an error

    let Some((_, items)) = content else {
        fatal!(ident.span() => "not a inline module"); // quit if the module is to a file
    };

    let mut item_store = SchemaStore::new();
    let mut output = HashMap::new();

    for item in items {
        let syn::Item::Struct(mut item_struct) = item else {
            item_store.insert(None, ModItem::Not(item.to_token_stream()));
            continue;
        };

        let attr: InnerAttr = match try_extract_attributes(&mut item_struct) {
            Ok(Some(attr)) => attr,
            Ok(None) => continue,
            Err(err) => {
                error!(err); // emit an error, but then continue with the rest
                continue;
            }
        };

        let ident = item_struct.ident.clone();
        let schema_item = SchemaItem { attr, item_struct };

        item_store.insert(Some(ident), ModItem::Schema(schema_item));
    }

    for item in &item_store {
        let (Some(ident), ModItem::Schema(schema_item)) = item else {
            continue;
        };

        output.insert(
            ident.clone(),
            match &schema_item.attr {
                InnerAttr::Row => TokenStream::new(),
                InnerAttr::Table(args) => table_impl(args, &schema_item.item_struct, &item_store),
            },
        );
    }

    let new_module: TokenStream = item_store
        .into_iter()
        .map(|(_, v)| match v {
            ModItem::Schema(schema_item) => {
                let original = schema_item.item_struct;
                let output = output.get(&original.ident);

                quote! {
                    #original

                    #output
                }
            }
            ModItem::Not(token_stream) => token_stream,
        })
        .collect();

    quote! {
        #(#attrs)*
        #vis
        #unsafety
        #mod_token
        #ident
        {
            #new_module
        }
        #semi
    }
}

// in a production proc macro, it would probably be better to use a custom implementation for compile time
type SchemaStore = ListOrderedMultimap<Option<syn::Ident>, ModItem>;

fn try_extract_attributes<T: deluxe::HasAttributes, R: deluxe::ExtractAttributes<T>>(
    obj: &mut T,
) -> deluxe::Result<Option<R>> {
    if obj.attrs().iter().any(|a| R::path_matches(a.path())) {
        return R::extract_attributes(obj).map(Some);
    }

    Ok(None)
}

enum ModItem {
    Schema(SchemaItem),
    Not(TokenStream),
}

struct SchemaItem {
    attr: InnerAttr,
    item_struct: syn::ItemStruct,
}

#[derive(deluxe::ExtractAttributes)]
#[deluxe(attributes(schema))]
enum InnerAttr {
    Row,
    #[deluxe(transparent)]
    Table(TableArgs),
}

#[derive(deluxe::ParseMetaItem)]
struct TableArgs {
    row: syn::Ident,
    pk: syn::Ident,
}

fn table_impl(args: &TableArgs, item_struct: &syn::ItemStruct, store: &SchemaStore) -> TokenStream {
    let ident = item_struct.ident.clone();
    let span = ident.span();

    let TableArgs { row, pk } = args;

    let getter_ident = format_ident!("find_by_{pk}");
    let Some(ModItem::Schema(row_item)) = store.get(&Some(row.clone())) else {
        bail!(row.span() => "row type not found in schema scope"); // emit an error and return the default
    };

    let syn::Fields::Named(row_fields) = &row_item.item_struct.fields else {
        bail!(row.span() => "row struct needs named fields"); // in this function, the default is an empty stream
    };

    let Some(pk_field) = row_fields
        .named
        .iter()
        .find(|f| f.ident.as_ref() == Some(pk))
    else {
        bail!(pk.span() => "primary key not found in row struct");
    };

    let getter_ty = &pk_field.ty;

    quote_spanned! {span=>
        impl #ident {
            #[allow(unused_variables)]
            pub fn #getter_ident(#pk: #getter_ty) -> #row {
                #row {
                    #pk: 32,
                    ..Default::default()
                }
            }
        }
    }
}