floop 0.3.0

use std::iter;

use unsynn::*;
use footgun_detector::footgun_detector;

use crate::{code_generation::footgun_detector::FootgunDetector, parser::{Arm, Ast, FuturePattern}};

mod footgun_detector;

fn hygenic_ident(str: &str) -> Ident {
    Ident::new(str, Span::mixed_site())
}

/// work around for `Ident::new` not being const.
macro_rules! declare_ident {
    ($ident:ident = $expr:expr; $($tokens:tt)+) => {
        declare_ident!($ident = $expr);
        declare_ident!($($tokens)*);
    };

    ($ident:ident = $expr:expr $(;)?) => {
        #[allow(non_camel_case_types, reason="the casing convays the type of the ident in the macro output")]
        struct $ident;

        impl ToTokens for $ident {
            fn to_tokens(&self, tokens: &mut TokenStream) {
                tokens.extend([hygenic_ident($expr)]);
            }
        }
    };
}

macro_rules! declare_indexed_ident {
    ($index_name:ident; $ident:ident = $expr:expr; $($tokens:tt)+) => {
        declare_indexed_ident!($index_name; $ident = $expr);
        declare_indexed_ident!($index_name; $($tokens)*);
    };

    ($index_name:ident; $ident:ident = $expr:expr $(;)?) => {
        #[allow(non_snake_case, reason="the casing convays the type of the ident in the macro output")]
        fn $ident($index_name: usize) -> Ident {
            hygenic_ident($expr)
        }
    };
}

// making the idents global is more conveniant, they're read-only anyways.
declare_ident! {
    result = "__FloopResult";
    unbiased_offset = "__floop_unbiased_offset";
    // this sadly isn't hygenic yet so more work arounds are used to prevent code from constructing the struct.
    DontBreak = "DontBreak";
    construct_dont_break = "__floop_construct_dont_break";
    index = "__floop_index";
}

declare_indexed_ident! {
    i;
    future = &format!("__floop_future_{i}");
    pin = &format!("__floop_pin_{i}");
    ResultVariant = &format!("Variant{i}");
    output = &format!("__floop_output{i}");
}

pub(crate) fn generate_code(ast: Ast) -> TokenStream {
    let before = ast.before.as_ref().map(|arm| arm.block()).unwrap_or_default();
    let after = ast.after.as_ref().map(|arm| arm.block()).unwrap_or_default();

    // the loop stops once all arms broke,
    // if there are no arms then all arms broke is always true (vacuous truth),
    // `before` still runs once as it runs before the check if all arms broke,
    // the rest of the loop can be removed.
    if ast.arms.is_empty() {
        return no_arms(before, after, ast.ret);
    }

    let mut tokens = TokenStream::new();
    let footgun_detector = footgun_detector();

    let mut poll_fn_arms = TokenStream::new();
    let mut match_arms = TokenStream::new();
    let mut recreate_futures = TokenStream::new();
    let mut break_condition = quote! {
        true
    };
    let mut break_value = TokenStream::new();
    let mut footgun_detectors = TokenStream::new();

    tokens.extend(gen_result(&ast));

    tokens.extend(quote! {
        struct #DontBreak;

        let #construct_dont_break = || #DontBreak;
        
        // ensure the async block never implements unpin.
        //i have no idea if this is actually necessary
        let __floop_marker = ::core::marker::PhantomPinned;
    });

    if !ast.biased {
        tokens.extend(quote! {
            let mut #unbiased_offset = 0;
        });
    }
    
    for (i, arm) in ast.arms.iter().enumerate() {
        let future = future(i);
        // this must be `Some` if `fut_name` is `Some` and `None` otherwise.
        let pin = pin(i);
        let output = output(i);
        let condition = arm.pattern().condition().map(|condition| quote! {
            && #condition
        }).unwrap_or(TokenStream::new());

        tokens.extend(quote! {
            let mut #future = ::core::option::Option::None;
            let mut #pin = ::core::option::Option::None;
            let mut #output = ::core::option::Option::None;
        });

        poll_fn_arms.extend(poll_fn_arm(i));

        recreate_futures.extend(recreate_future(i, arm, &condition));

        match_arms.extend(match_arm(i, arm));

        break_condition.extend(quote! {
            && let Some(#output) = #output
        });

        // `(expr,)` is the syntax for single element tuples, which i didn't even know existed,
        // so trailing commas must not be generated as they would make the break value of a single-arm floop inconvenient to use.
        let comma = (i != 0).then_some(quote! { , });
        break_value.extend(quote! {
            #comma #output
        });

        if !arm.pattern().footgun_allowed() {
            footgun_detectors.extend(gen_footgun_detectors(i, arm, &footgun_detector, condition));
        }
    }

    tokens.extend(combine_fragments(&ast, before, after, poll_fn_arms, match_arms, recreate_futures, break_condition, break_value));

    let footgun_declerations = footgun_detector.declerations;
    let value = hygenic_ident("__floop_value");

    if ast.ret {
        finish_with_return(tokens, footgun_declerations, footgun_detectors)
    } else {
        #[expect(non_snake_case)]
        let DontReturn = hygenic_ident("__FloopDontReturn");
        let construct_dont_return = hygenic_ident("__floop_construct_dont_return");
        // ensure the macro can be used as a expression.
        quote! {
            async {
                struct #DontReturn;

                let #construct_dont_return = || #DontReturn;

                let (#value, _): (_, #DontReturn) = async {
                    // the non-hygenic-struct workaround again.
                    struct #DontReturn;

                    let #value = {
                        #tokens
                    };
                    
                    (#value, #construct_dont_return())
                }.await;

                // the detectors need to be at the end for let conditions creating futures whose type is inferred based on a value produced inside a arm to compile.
                {
                    #footgun_declerations

                    #footgun_detectors;
                }

                #value
            }
        }
    }
}

fn finish_with_return(tokens: TokenStream, footgun_declerations: TokenStream, footgun_detectors: TokenStream) -> TokenStream {

    let run = hygenic_ident("__floop_run");
    let value = hygenic_ident("__floop_value");
    // TODO: this needs more tests.
    quote! {
        {
            let #value = async {
                let mut #value = ::core::option::Option::None;
                // FIXME: i am not sure which ordering is correct so i chose `SeqCst`, a weaker ordering should be chosen if possible.
                //     somehow getting rid of the atomic entirely would be even better, but that seems very difficult.
                // `value` can't be read to determine whether or not the future is done while the future is running
                // as `value` is mutably borrowed until the future is canceled, which creates a chicken or egg problem.
                let #run = ::core::sync::atomic::AtomicBool::new(true);

                // this horrible mess of spaghetti just catches a `return` and returns from the outer function instead,
                // the control flow is *worse* than it looks.
                // i hope there's a better way i'm missing.
                // i wonder how badly this optimizes.
                {
                    let mut workaround = ::core::pin::pin!(async {
                        // i'm not sure if i can explain this code clearly, but it works roughly like this:
                        // if the arms don't return then the outer `=` never runs, but the inner `=` does, resulting in `value` being set to continue,
                        // the atomic stops the future as it would otherwise wait forever,
                        // if the arm does return the inner `=` never runs, but the outer `=` does,
                        // and the atomic prevents the future from being polled after it finished/
                        #value = ::core::option::Option::Some(::core::ops::ControlFlow::Break(async {
                            #value = ::core::option::Option::Some(::core::ops::ControlFlow::Continue({
                                #tokens
                            }));

                            ::core::sync::atomic::AtomicBool::store(&#run, false, ::core::sync::atomic::Ordering::SeqCst);

                            // `yield_now` is not in core so i need to reimplement it.
                            // ensure `run` always has a oppertunity to be read the atomic.
                            let mut yielded = false;
                            ::core::future::poll_fn(|ctx| {
                                if !yielded {
                                    yielded = true;
                                    
                                    let waker = ::core::task::Context::waker(ctx);
                                    ::core::task::Waker::wake_by_ref(waker);

                                    ::core::task::Poll::Pending
                                } else {
                                    ::core::task::Poll::Ready(())
                                }
                            }).await;

                            ::core::unreachable!()
                        }.await));

                        ::core::sync::atomic::AtomicBool::store(&#run, false, ::core::sync::atomic::Ordering::SeqCst);
                    });

                    let poll_fn = ::core::future::poll_fn(|ctx| {
                        if ::core::sync::atomic::AtomicBool::load(&#run, ::core::sync::atomic::Ordering::SeqCst) {
                            workaround.as_mut().poll(ctx)
                        } else {
                            ::core::task::Poll::Ready(())
                        }
                    });
                    poll_fn.await;
                }

                // the detectors need to be at the end for let conditions creating futures whose type is inferred based on a value produced inside a arm to compile.
                {
                    #footgun_declerations

                    #footgun_detectors;
                }

                ::core::option::Option::unwrap(#value)
            }.await;

            match #value {
                ::core::ops::ControlFlow::Continue(con) => con,
                ::core::ops::ControlFlow::Break(brk) => return brk,
            }
        }
    }
}

fn gen_result(ast: &Ast) -> TokenStream {
    let tys = (0..ast.arms.len()).map(|i| hygenic_ident(&format!("T{i}")));
    let variants = tys.clone().enumerate().map(|(i, ty)| (ResultVariant(i), ParenthesisGroupContaining::new(ty), Punct::new(',', Spacing::Alone)));
    let tys = tys.zip(iter::repeat(Punct::new(',', Spacing::Alone)));

    quote! {
        enum #result<#{tys}> {
            #{variants}
        }
    }
}

#[expect(clippy::too_many_arguments)]
fn combine_fragments(
    ast: &Ast,
    before: TokenStream,
    after: TokenStream,
    poll_fn_arms: TokenStream,
    match_arms: TokenStream,
    recreate_futures: TokenStream,
    break_condition: TokenStream,
    break_value: TokenStream
) -> TokenStream{
    let poll_fn = hygenic_ident("__floop_poll_fn");
    let arm_count = ast.arms.len();
    let unbias = (!ast.biased).then_some(unbias(ast));

    quote! {
        let tmp: (#DontBreak, _) = loop {
            struct #DontBreak;
        
            #before;
        
            #recreate_futures
        
            let #poll_fn = |ctx: &mut ::core::task::Context| {
                for #index in 0..#arm_count {
                    #unbias

                    match #index {
                        #poll_fn_arms
                        // all numbers between `0` and `arm_amount` are matched by a arm's branch.
                        _ => ::core::unreachable!(),
                    }
                }

                ::core::task::Poll::Pending
            };
            let #poll_fn = ::core::future::poll_fn(#poll_fn);

            match #poll_fn.await {
                #match_arms
            }

            if #break_condition {
                break (#construct_dont_break(), (#break_value));
            }

            #after;
        };

        tmp.1
    }
}

fn unbias(ast: &Ast) -> TokenStream {
    let arm_count = ast.arms.len();
    quote! {
        // 0 arms result in a early return, so this can't panic.
        #[allow(clippy::modulo_one)]
        let #index = (#index + #unbiased_offset) % #arm_count;
        #unbiased_offset += 1;
        #[allow(clippy::modulo_one)]
        {
            #unbiased_offset %= #arm_count
        };
    }
}

fn no_arms(before: TokenStream, after: TokenStream, ret: bool) -> TokenStream {
    let await_ = ret.then_some(quote! {
        .await
    });

    quote! {
        {
            async {
                loop {
                    #before;
                    break;
                    // ensure invalid `after` code still results in a compiler error.
                    #after;
                };
            }#await_
        }
    }
}

fn gen_footgun_detectors(i: usize, arm: &Arm<FuturePattern>, footgun_detector: &FootgunDetector, condition: TokenStream) -> TokenStream {
    let future = future(i);
    let expr = arm.pattern().expr();
    let detector = &footgun_detector.detector;

    quote! {
        // the closure bypasses a bug in the compiler (const panics in async blocks sometimes compile) and
        // ensures that the detector doesn't actually do anything at runtime, as the closure is never called.
         let _ = || {
            if true #condition {
                let mut #future = Some({#expr});
                #future = #detector;
            }
        };
    }
}

fn match_arm(i: usize, arm: &Arm<FuturePattern>) -> TokenStream {
    let future = future(i);
    let pin = pin(i);
    let result_variant = ResultVariant(i);
    let output = output(i);
    let block = &arm.block();
    let pattern = &arm.pattern().pattern();
    let output_storage = hygenic_ident("__floop_output_storage");
    let matched = hygenic_ident("__floop_matched");
    let second_iteration = hygenic_ident("__floop_second_iteration");
    let should_continue = hygenic_ident("__floop_should_continue");

    quote! {
        // TODO: add the ability to use `break` to break specific arms.
        // (and the entire loop should stop once all arms broke).
        #result::#result_variant(#matched) => {
            let #pattern = #matched;

            let mut #should_continue = false;
            
            // this weird mess of loops catches `break` and breaks the specific arm instead of the entire loop.
            // it also catches continue and continues the loop after setting the pin and future to `None`.
            // FIXME: this allow likely also applies to `block`.
            #[allow(unreachable_code, clippy::never_loop, unused_variables)]
            '__floop_outer: loop {
                let mut #second_iteration = false;
                let #output_storage = loop {
                    if #second_iteration {
                        #should_continue = true;

                        break '__floop_outer;
                    }
                    #second_iteration = true;
                    
                    #block;

                    break '__floop_outer;
                };

                #output = ::core::option::Option::Some(#output_storage);
                break;
            }

            // the future needs to be recreated so that it can be polled again.
            #pin = ::core::option::Option::None;
            // SAFETY: assignment runs `drop` before overwriting the value so this is safe,
            // the safety contract of `Pin` is now over as the value is dropped.
            #future = ::core::option::Option::None;

            /// the loop can only be continued after the future and pin where set to none,
            /// otherwise the future will be polled after it finished.
            if #should_continue {
                continue;
            }
        }
    }
}

fn recreate_future(i: usize, arm: &Arm<FuturePattern>, condition: &TokenStream) -> TokenStream {
    let future = future(i);
    let pin = pin(i);
    let output = output(i);
    let fut_expr = &arm.pattern().expr();
    let tmp = hygenic_ident("__floop_fut_ref");
    let future_ref = hygenic_ident("__floop_future_ref");

    quote! {
        // `fut_name` can't be used because it is borrowed by `pin_name`.
        if ::core::option::Option::is_none(&#pin) && ::core::option::Option::is_none(&#output) #condition {
            // this shouldn't be needed, but it shouldn't hurt.
            #pin = ::core::option::Option::None;
            // this may help the compiler undertsand that `fut_name` is `None`, i have no idea if it actually does anything.
            #future = ::core::option::Option::None;
            // FIXME: `fut_expr` can still break, this only prevents break from doing anything, not from being called.
            let #tmp = loop {
                break {
                    #fut_expr
                }
            };
        
            let #future_ref = ::core::option::Option::get_or_insert(&mut #future, {
                #tmp
            });

        
            // SAFETY: `fut_name` will not be moved and it will be dropped before it's invalidated,
            // `pinned` borrows it and both are always set to none directly after each other,
            // so there is no point where unrelated code runs and moving `fut_name` doesn't result in a compiler error.
            // this starts a new `Pin` safety contract.
            // `future` and `pin` is `mixed_site` it's impossible for user code to change either.
            // (which would lead to undefined behaviour if it was possible).
            unsafe {
                #pin = ::core::option::Option::Some(::core::pin::Pin::new_unchecked(#future_ref));
            }
        }
    }
}

fn poll_fn_arm(i: usize) -> TokenStream {
    let pin = pin(i);
    let result_variant = ResultVariant(i);
    
    quote! {
        #i => {
            if let Some(ref mut #pin) = #pin {
                let #pin = ::core::pin::Pin::as_mut(#pin);
                // this consumes `pin`, preventing `fut_expr` from using it.
                if let ::core::task::Poll::Ready(output) = ::core::future::Future::poll(#pin, ctx) {
                    return ::core::task::Poll::Ready(#result::#result_variant(output));
                }
            }
        }
    }
}