// Copyright 2020 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

extern crate proc_macro;

use {
    proc_macro::TokenStream,
    quote::{quote, quote_spanned},
    syn::{
        parse,
        parse::{Parse, ParseStream},
        Attribute, Block, Error, Ident, ItemFn, LitBool, LitInt, LitStr, Signature, Token,
    },
};

#[derive(Clone, Copy)]
enum FunctionType {
    Component,
    Test,
}

// How should code be executed?
#[derive(Clone, Copy)]
enum Executor {
    // Directly by calling it
    None,
    // fasync::run_singlethreaded
    Singlethreaded,
    // fasync::run
    Multithreaded(usize),
    // #[test]
    Test,
    // fasync::run_singlethreaded(test)
    SinglethreadedTest,
    // fasync::run(test)
    MultithreadedTest(usize),
    // fasync::run_until_stalled
    UntilStalledTest,
}

impl Executor {
    fn is_test(&self) -> bool {
        match self {
            Executor::Test
            | Executor::SinglethreadedTest
            | Executor::MultithreadedTest(_)
            | Executor::UntilStalledTest => true,
            Executor::None | Executor::Singlethreaded | Executor::Multithreaded(_) => false,
        }
    }
}

impl quote::ToTokens for Executor {
    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
        tokens.extend(match self {
            Executor::None => quote! { ::fuchsia::main_not_async(func) },
            Executor::Test => quote! { ::fuchsia::test_not_async(func) },
            Executor::Singlethreaded => quote! { ::fuchsia::main_singlethreaded(func) },
            Executor::Multithreaded(n) => quote! { ::fuchsia::main_multithreaded(func, #n) },
            Executor::SinglethreadedTest => quote! { ::fuchsia::test_singlethreaded(func) },
            Executor::MultithreadedTest(n) => quote! { ::fuchsia::test_multithreaded(func, #n) },
            Executor::UntilStalledTest => quote! { ::fuchsia::test_until_stalled(func) },
        })
    }
}

// Helper trait for things that can generate the final token stream
trait Finish {
    fn finish(self) -> TokenStream
    where
        Self: Sized;
}

struct Transformer {
    executor: Executor,
    attrs: Vec<Attribute>,
    sig: Signature,
    block: Box<Block>,
    logging: bool,
}

struct Args {
    threads: usize,
    allow_stalls: bool,
    logging: bool,
}

fn get_arg<T: Parse>(p: &ParseStream<'_>) -> syn::Result<T> {
    p.parse::<Token![=]>()?;
    p.parse()
}

fn get_base10_arg<T>(p: &ParseStream<'_>) -> syn::Result<T>
where
    T: std::str::FromStr,
    T::Err: std::fmt::Display,
{
    get_arg::<LitInt>(p)?.base10_parse()
}

fn get_bool_arg(p: &ParseStream<'_>, if_present: bool) -> syn::Result<bool> {
    if p.peek(Token![=]) {
        Ok(get_arg::<LitBool>(p)?.value)
    } else {
        Ok(if_present)
    }
}

impl Parse for Args {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let mut args = Self { threads: 1, allow_stalls: true, logging: true };

        loop {
            if input.is_empty() {
                break;
            }
            let ident: Ident = input.parse()?;
            let err = |message| Err(Error::new(ident.span(), message));
            match ident.to_string().as_ref() {
                "threads" => args.threads = get_base10_arg(&input)?,
                "allow_stalls" => args.allow_stalls = get_bool_arg(&input, true)?,
                "logging" => args.logging = get_bool_arg(&input, true)?,
                x => return err(format!("unknown argument: {}", x)),
            }
            if input.is_empty() {
                break;
            }
            input.parse::<Token![,]>()?;
        }

        Ok(args)
    }
}

impl Transformer {
    // Construct a new Transformer, verifying correctness.
    fn parse(
        function_type: FunctionType,
        args: TokenStream,
        input: TokenStream,
    ) -> Result<Transformer, Error> {
        let args: Args = parse(args)?;
        let ItemFn { attrs, vis: _, sig, block } = parse(input)?;
        let is_async = sig.asyncness.is_some();

        let err = |message| Err(Error::new(sig.ident.span(), message));

        let executor = match (args.threads, args.allow_stalls, is_async, function_type) {
            (0, _, _, _) => return err("need at least one thread"),
            (_, false, _, FunctionType::Component) => {
                return err("allow_stalls=false only applies to tests")
            }
            (1, _, false, FunctionType::Component) => Executor::None,
            (1, true, true, FunctionType::Component) => Executor::Singlethreaded,
            (n, true, true, FunctionType::Component) => Executor::Multithreaded(n),
            (1, _, false, FunctionType::Test) => Executor::Test,
            (1, true, true, FunctionType::Test) => Executor::SinglethreadedTest,
            (n, true, true, FunctionType::Test) => Executor::MultithreadedTest(n),
            (1, false, true, FunctionType::Test) => Executor::UntilStalledTest,
            (_, false, _, FunctionType::Test) => {
                return err("allow_stalls=false tests must be single threaded")
            }
            (_, true, false, _) => return err("must be async to use >1 thread"),
        };

        Ok(Transformer { executor, attrs, sig, block, logging: args.logging })
    }
}

impl Finish for Transformer {
    // Build the transformed code, knowing that everything is ok because we proved that in parse.
    fn finish(self) -> TokenStream {
        let ident = self.sig.ident;
        let span = ident.span();
        let ret_type = self.sig.output;
        let attrs = self.attrs;
        let asyncness = self.sig.asyncness;
        let block = self.block;
        let inputs = self.sig.inputs;

        let mut func_attrs = Vec::new();

        // Initialize logging
        let init_logging = if !self.logging {
            quote! {}
        } else if self.executor.is_test() {
            let test_name = LitStr::new(&format!("{}", ident), ident.span());
            quote! { ::fuchsia::init_logging_for_test(#test_name); }
        } else {
            quote! { ::fuchsia::init_logging_for_component(); }
        };

        if self.executor.is_test() {
            // Add test attribute to outer function.
            func_attrs.push(quote!(#[test]));
        }

        let func = if self.executor.is_test() {
            quote! { test_entry_point }
        } else {
            quote! { component_entry_point }
        };

        // Adapt the runner function based on whether it's a test and argument count
        // by providing needed arguments.
        let adapt_main = match (self.executor.is_test(), inputs.len()) {
            // Main function, no arguments - no adaption needed.
            (false, 0) => quote! { #func },
            // Main function, one arguemnt - adapt by parsing command line arguments.
            (false, 1) => quote! { ::fuchsia::adapt_to_parse_arguments(#func) },
            // Test function, no arguments - adapt by taking the run number and discarding it.
            (true, 0) => quote! { ::fuchsia::adapt_to_take_test_run_number(#func) },
            // Test function, one argument - no adaption needed.
            (true, 1) => quote! { #func },
            // Anything with more than one argument: error.
            (_, n) => panic!("Too many ({}) arguments to function", n),
        };

        // Select executor
        let run_executor = self.executor;

        // Finally build output.
        let output = quote_spanned! {span =>
            #(#attrs)* #(#func_attrs)*
            fn #ident () #ret_type {
                #asyncness fn #func(#inputs) #ret_type {
                    #block
                }
                #init_logging
                let func = #adapt_main;
                #run_executor
            }
        };
        output.into()
    }
}

impl Finish for Error {
    fn finish(self) -> TokenStream {
        self.to_compile_error().into()
    }
}

impl<R: Finish, E: Finish> Finish for Result<R, E> {
    fn finish(self) -> TokenStream {
        match self {
            Ok(r) => r.finish(),
            Err(e) => e.finish(),
        }
    }
}

/// Define a fuchsia component.
///
/// This attribute should be applied to the process `main` function.
/// It will take care of setting up various Fuchsia crates for the component.
/// If an async function is provided, a fuchsia-async Executor will be used to execute it.
///
/// Arguments:
///  - `threads` - integer worker thread count for the component. Must be >0. Default 1.
///  - `logging` - boolean toggle for whether to initialize logging (or not). Default true.
///                This currently does nothing on host. On Fuchsia fuchsia-syslog is used.
///
/// The main function can return either () or a Result<(), E> where E is an error type.
/// If the main function takes an argument, it's expected that argument implement
/// argh::TopLevelCommand. argh::from_env will be invoked to parse command line arguments and
/// pass them to the main function.
#[proc_macro_attribute]
pub fn component(args: TokenStream, input: TokenStream) -> TokenStream {
    Transformer::parse(FunctionType::Component, args, input).finish()
}

/// Define a fuchsia test.
///
/// This attribute should be applied to a function in a cfg(test) module.
/// It will take care of setting up various Fuchsia crates for the test.
/// If an async function is provided, a fuchsia-async Executor will be used to execute it.
///
/// Arguments:
///  - `threads`      - integer worker thread count for the test. Must be >0. Default 1.
///  - `logging`      - boolean toggle for whether to initialize logging (or not). Default true.
///                     This currently does nothing on host. On Fuchsia fuchsia-syslog is used.
///  - `allow_stalls` - boolean toggle for whether the async test is allowed to stall during
///                     execution (if true), or whether the function must complete without pausing
///                     (if false).
///                     `.await` is not a stall if something preceding the await will guarantee
///                     that it finishes within one loop of the Executor. Defaults to true.
///                     This argument is not currently available for host tests.
///
/// The test function can return either () or a Result<(), E> where E is an error type.
/// The test function can either take no arguments, or a single usize argument. If it takes an
/// argument, that value will be the current iteration count of running this test repeatedly.
#[proc_macro_attribute]
pub fn test(args: TokenStream, input: TokenStream) -> TokenStream {
    Transformer::parse(FunctionType::Test, args, input).finish()
}