use proc_macro2::TokenStream;
use quote::quote;
use syn::{parse2, ItemFn, Type, Token};
use syn::parse::{Parse, ParseStream};

/// Parsed attribute for `#[mad_hatter::main(TxType)]` or `#[mad_hatter::main(TxType @ DB)]`.
struct MainAttr {
    tx_type: Type,
    db_binding: Option<syn::Path>,
}

impl Parse for MainAttr {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let tx_type: Type = input.parse()?;
        let db_binding = if input.peek(Token![@]) {
            input.parse::<Token![@]>()?;
            Some(input.parse::<syn::Path>()?)
        } else {
            None
        };
        Ok(MainAttr { tx_type, db_binding })
    }
}

/// Expand `#[mad_hatter::main(TxType)]` or `#[mad_hatter::main(TxType @ DB)]` attribute macro.
///
/// Transforms:
/// ```ignore
/// #[mad_hatter::main(AppStartup @ MEMORY_DB)]
/// async fn main(scope: &TxScope<AppStartup>) {
///     // user body
/// }
/// ```
///
/// Into:
/// ```ignore
/// fn main() {
///     mad_hatter::__tokio::runtime::Builder::new_current_thread()
///         .enable_all()
///         .build()
///         .unwrap()
///         .block_on(async {
///             mad_hatter::__main_entry::<AppStartup, _>(
///                 Some(MEMORY_DB.get().expect("DB not initialized").db().clone()),
///                 |__mad_scope| async move {
///                     let scope = &*__mad_scope;
///                     // user body
///                     mad_hatter::__start_framework().await;
///                 },
///             ).await;
///         })
/// }
/// ```
pub fn expand_main(attr: TokenStream, input: TokenStream) -> TokenStream {
    // Parse Transaction type and optional DB binding from attribute
    let main_attr: MainAttr = parse2(attr)
        .expect("#[mad_hatter::main] requires a Transaction type argument, e.g. #[mad_hatter::main(AppStartup)] or #[mad_hatter::main(AppStartup @ DB)]");

    let tx_type = &main_attr.tx_type;

    // Parse the async fn
    let func: ItemFn = parse2(input)
        .expect("#[mad_hatter::main] must be applied to an async fn");

    // Validate: must be async
    if func.sig.asyncness.is_none() {
        panic!("#[mad_hatter::main] function must be async");
    }

    // Validate: function name must be "main"
    if func.sig.ident != "main" {
        panic!("#[mad_hatter::main] must be applied to a function named `main`");
    }

    // Extract scope parameter name from first (and only) parameter
    let param_name = match func.sig.inputs.first() {
        Some(syn::FnArg::Typed(pat_type)) => {
            match &*pat_type.pat {
                syn::Pat::Ident(pat_ident) => pat_ident.ident.clone(),
                _ => panic!("#[mad_hatter::main] scope parameter must be a simple identifier"),
            }
        }
        _ => panic!("#[mad_hatter::main] requires exactly one parameter: scope: &TxScope<TxType>"),
    };

    // Validate: exactly one parameter
    if func.sig.inputs.len() != 1 {
        panic!("#[mad_hatter::main] requires exactly one parameter (scope)");
    }

    // Extract function body statements
    let stmts = &func.block.stmts;

    // Generate db expression
    let db_expr = if let Some(ref db_path) = main_attr.db_binding {
        quote! { Some(#db_path.get().expect("DB not initialized").db().clone()) }
    } else {
        quote! { None }
    };

    quote! {
        fn main() {
            mad_hatter::__tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap()
                .block_on(async {
                    mad_hatter::__main_entry::<#tx_type, _>(#db_expr, |__mad_scope| async move {
                        let #param_name = &*__mad_scope;
                        #(#stmts)*
                        mad_hatter::__start_framework().await;
                    }).await;
                })
        }
    }
}
