// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//

#pragma once

#include <boost/asio/executor_work_guard.hpp>
#include <boost/asio/io_context.hpp>
#include <boost/asio/spawn.hpp>
#include <boost/optional.hpp>
#include <boost/thread/csbl/memory/allocator_arg.hpp>

#include <condition_variable>
#include <cstddef>
#include <exception>
#include <mutex>
#include <thread>
#include <vector>

namespace beast::test {

/**
 * Mix-in to support tests using asio coroutines.
 *
 * Derive from this class and use yield_to to launch test
 * functions inside coroutines. This is handy for testing
 * asynchronous asio code.
 */
class EnableYieldTo
{
protected:
    boost::asio::io_context ios_;

private:
    boost::optional<boost::asio::executor_work_guard<boost::asio::io_context::executor_type>> work_;
    std::vector<std::thread> threads_;
    std::mutex m_;
    std::condition_variable cv_;
    std::size_t running_ = 0;

public:
    /**
     * The type of yield context passed to functions.
     */
    using yield_context = boost::asio::yield_context;

    explicit EnableYieldTo(std::size_t concurrency = 1) : work_(boost::asio::make_work_guard(ios_))
    {
        threads_.reserve(concurrency);
        for (std::size_t i = 0; i < concurrency; ++i)
        {
            threads_.emplace_back([&] { ios_.run(); });
        }
    }

    ~EnableYieldTo()
    {
        work_ = boost::none;
        for (auto& t : threads_)
            t.join();
    }

    /**
     * Return the `io_context` associated with the object
     */
    boost::asio::io_context&
    getIoContext()
    {
        return ios_;
    }

    /**
     * Run one or more functions, each in a coroutine.
     *
     * This call will block until all coroutines terminate.
     *
     * Each functions should have this signature:
     * @code
     *     void f(yield_context);
     * @endcode
     *
     * @param fn... One or more functions to invoke.
     */
#if BEAST_DOXYGEN
    template <class... FN>
    void
    yield_to(FN&&... fn);
#else
    template <class F0, class... FN>
    void
    yieldTo(F0&& f0, FN&&... fn);
#endif

private:
    void
    spawn()
    {
    }

    template <class F0, class... FN>
    void
    spawn(F0&& f, FN&&... fn);
};

template <class F0, class... FN>
void
EnableYieldTo::yieldTo(F0&& f0, FN&&... fn)
{
    running_ = 1 + sizeof...(FN);
    spawn(f0, fn...);
    std::unique_lock<std::mutex> lock{m_};
    cv_.wait(lock, [&] { return running_ == 0; });
}

template <class F0, class... FN>
inline void
EnableYieldTo::spawn(F0&& f, FN&&... fn)
{
    boost::asio::spawn(
        ios_,
        boost::allocator_arg,
        boost::context::fixedsize_stack(2 * 1024 * 1024),
        [&](yield_context yield) {
            f(yield);
            std::scoped_lock const lock{m_};
            if (--running_ == 0)
                cv_.notify_all();
        },
        [](std::exception_ptr e) {
            if (e)
                std::rethrow_exception(e);
        });
    spawn(fn...);
}

}  // namespace beast::test