//
// A footgun that survives "doing everything right".
//
// This program defines a custom coroutine return type, my_task<Eager, T>, that:
//
//   * satisfies capy::IoAwaitable (there is a static_assert below proving it),
//   * implements the two-argument await_suspend(handle, io_env const*) and
//     faithfully stores the environment via set_environment(),
//   * provides a transform_awaitable that injects the environment into nested
//     IoAwaitables, exactly like capy::task.
//
// In other words, the type follows the IoAwaitable protocol correctly.
//
// The ONLY difference between the safe path and the undefined-behavior path is
// a single word in initial_suspend():
//
//     Eager == false  ->  std::suspend_always   (lazy: env installed first)
//     Eager == true   ->  std::suspend_never     (body runs before env exists)
//
// The coroutine body then does something entirely reasonable: run a little
// setup, then check the stop token before starting expensive work. With the
// lazy variant this works. With the eager variant the body runs before
// await_suspend was ever called, so the promise's env_ is still nullptr, and
//
//     co_await capy::this_coro::stop_token
//
// dereferences it inside io_awaitable_promise_base::await_transform:
//
//   * Debug build (assertions on): BOOST_CAPY_ASSERT(env_) fires -> abort.
//   * Release build (NDEBUG):      env_->stop_token dereferences null -> UB.
//
// Nothing in the call site looks wrong: app() is an ordinary capy::task that
// simply `co_await`s the child, launched with run_async on a thread_pool.
//

#include <boost/capy.hpp>
#include <boost/capy/concept/io_awaitable.hpp>
#include <boost/capy/ex/io_awaitable_promise_base.hpp>
#include <boost/capy/ex/io_env.hpp>
#include <boost/capy/ex/this_coro.hpp>

#include <coroutine>
#include <exception>
#include <iostream>
#include <optional>
#include <type_traits>
#include <utility>

namespace capy = boost::capy;

namespace footgun {

//--------------------------------------------------------------------
// Result storage for my_task<Eager, T>.
//--------------------------------------------------------------------
template<class T>
struct task_result
{
    std::optional<T> value_;
    std::exception_ptr err_;
    void return_value(T v) { value_.emplace(std::move(v)); }
    void unhandled_exception() noexcept { err_ = std::current_exception(); }
    std::exception_ptr exception() const noexcept { return err_; }
    T& result() { return *value_; }
};

template<>
struct task_result<void>
{
    std::exception_ptr err_;
    void return_void() noexcept {}
    void unhandled_exception() noexcept { err_ = std::current_exception(); }
    std::exception_ptr exception() const noexcept { return err_; }
};

//--------------------------------------------------------------------
// A protocol-conforming custom IoAwaitable coroutine type.
//
// The only knob is `Eager`, which selects initial_suspend(). Everything else
// is the textbook implementation.
//--------------------------------------------------------------------
template<bool Eager, class T = void>
class my_task
{
public:
    struct promise_type
        : capy::io_awaitable_promise_base<promise_type>
        , task_result<T>
    {
        my_task get_return_object()
        {
            return my_task{
                std::coroutine_handle<promise_type>::from_promise(*this)};
        }

        // THE knob. Both are legal; only the timing of the body differs.
        auto initial_suspend() noexcept
        {
            if constexpr(Eager)
                return std::suspend_never{};  // body runs at construction
            else
                return std::suspend_always{}; // body waits until resumed
        }

        auto final_suspend() noexcept
        {
            struct awaiter
            {
                promise_type* p;
                bool await_ready() const noexcept { return false; }
                std::coroutine_handle<>
                await_suspend(std::coroutine_handle<>) const noexcept
                {
                    return p->continuation();
                }
                void await_resume() const noexcept {}
            };
            return awaiter{this};
        }

        // Faithful protocol behavior: inject env into nested IoAwaitables,
        // just like capy::task does.
        template<class A>
        struct transform_awaiter
        {
            std::decay_t<A> a_;
            promise_type* p_;
            bool await_ready() { return a_.await_ready(); }
            auto await_suspend(std::coroutine_handle<> h)
            {
                return a_.await_suspend(h, p_->environment());
            }
            decltype(auto) await_resume() { return a_.await_resume(); }
        };

        template<class A>
        auto transform_awaitable(A&& a)
        {
            static_assert(capy::IoAwaitable<std::decay_t<A>>,
                "co_await inside my_task requires an IoAwaitable");
            return transform_awaiter<A>{std::forward<A>(a), this};
        }
    };

    my_task(my_task&& o) noexcept : h_(std::exchange(o.h_, {})) {}
    my_task(my_task const&) = delete;
    ~my_task() { if(h_) h_.destroy(); }

    //---- IoAwaitable interface: the "I followed the protocol" part ----

    bool await_ready() const noexcept { return false; }

    std::coroutine_handle<>
    await_suspend(std::coroutine_handle<> cont, capy::io_env const* env)
    {
        h_.promise().set_continuation(cont);
        h_.promise().set_environment(env); // stored faithfully, as required
        return h_;
    }

    T await_resume()
    {
        auto& p = h_.promise();
        if(p.exception())
            std::rethrow_exception(p.exception());
        if constexpr(!std::is_void_v<T>)
            return std::move(p.result());
    }

    //---- IoRunnable bits ----
    std::coroutine_handle<promise_type> handle() const noexcept { return h_; }
    void release() noexcept { h_ = {}; }

private:
    explicit my_task(std::coroutine_handle<promise_type> h) noexcept : h_(h) {}
    std::coroutine_handle<promise_type> h_;
};

// Proof that the type follows the protocol: both variants ARE IoAwaitables.
static_assert(capy::IoAwaitable<my_task<false, int>>,
    "lazy my_task must satisfy IoAwaitable");
static_assert(capy::IoAwaitable<my_task<true, int>>,
    "eager my_task must satisfy IoAwaitable");

//--------------------------------------------------------------------
// A perfectly reasonable coroutine body.
//--------------------------------------------------------------------
inline int
expensive_setup()
{
    std::cout << "    [child] running setup work in the eager prefix...\n";
    return 41;
}

template<bool Eager>
my_task<Eager, int>
compute()
{
    // ---- body that runs before the first real suspension ----
    int partial = expensive_setup();

    // Looks completely reasonable: check for cancellation before doing the
    // expensive I/O below. For Eager==true this executes before await_suspend
    // ever ran, so the promise env_ is still nullptr.
    auto token = co_await capy::this_coro::stop_token; // Eager -> null deref / UB
    if(token.stop_requested())
        co_return 0;

    // (Real code would co_await actual I/O here.)
    co_return partial + 1;
}

} // namespace footgun

//--------------------------------------------------------------------
// An ordinary-looking application coroutine. Nothing here hints at trouble.
//--------------------------------------------------------------------
capy::task<>
app()
{
    // 1) Lazy variant: identical body, identical protocol -> works.
    std::cout << "[app] awaiting LAZY compute()...\n";
    int a = co_await footgun::compute<false>();
    std::cout << "[app] lazy result = " << a << "\n\n";

    // 2) Eager variant: only initial_suspend() differs -> UB while the
    //    operand below is being constructed (its body runs immediately).
    std::cout << "[app] awaiting EAGER compute()...\n";
    std::cout.flush();
    int b = co_await footgun::compute<true>();
    std::cout << "[app] eager result = " << b << " (never reached)\n";
    co_return;
}

int
main()
{
    capy::thread_pool pool;
    capy::run_async(pool.get_executor())(app());
    pool.join();
    return 0;
}