#include "util/CoroutineGroup.hpp"

#include "util/Assert.hpp"
#include "util/Spawn.hpp"

#include <boost/asio/spawn.hpp>
#include <boost/asio/steady_timer.hpp>

#include <cstddef>
#include <functional>
#include <optional>
#include <utility>

namespace util {

CoroutineGroup::CoroutineGroup(boost::asio::yield_context yield, std::optional<size_t> maxChildren)
    : timer_{yield.get_executor(), boost::asio::steady_timer::duration::max()}
    , maxChildren_{maxChildren}
{
}

CoroutineGroup::~CoroutineGroup()
{
    ASSERT(
        childrenCounter_ == 0,
        "CoroutineGroup is destroyed without waiting for child coroutines to finish"
    );
}

bool
CoroutineGroup::spawn(
    boost::asio::yield_context yield,
    std::function<void(boost::asio::yield_context)> fn
)
{
    if (isFull())
        return false;

    ++childrenCounter_;
    util::spawn(yield, [this, fn = std::move(fn)](boost::asio::yield_context yield) {
        fn(yield);
        onCoroutineCompleted();
    });
    return true;
}

std::optional<std::function<void()>>
CoroutineGroup::registerForeign(boost::asio::yield_context yield)
{
    if (isFull())
        return std::nullopt;

    ++childrenCounter_;
    // It is important to spawn onCoroutineCompleted() to the same coroutine as will be calling
    // asyncWait(). timer_ here is not thread safe, so without spawn there could be a data race.
    return [this, yield]() { util::spawn(yield, [this](auto&&) { onCoroutineCompleted(); }); };
}

void
CoroutineGroup::asyncWait(boost::asio::yield_context yield)
{
    if (childrenCounter_ == 0)
        return;

    boost::system::error_code error;
    timer_.async_wait(yield[error]);
}

size_t
CoroutineGroup::size() const
{
    return childrenCounter_;
}

bool
CoroutineGroup::isFull() const
{
    return maxChildren_.has_value() && childrenCounter_ >= *maxChildren_;
}

void
CoroutineGroup::onCoroutineCompleted()
{
    ASSERT(
        childrenCounter_ != 0,
        "onCoroutineCompleted() called more times than the number of child coroutines"
    );

    --childrenCounter_;
    if (childrenCounter_ == 0)
        timer_.cancel();
}

}  // namespace util