#include "util/AsioContextTestFixture.hpp"
#include "util/CoroutineGroup.hpp"
#include "util/Spawn.hpp"

#include <boost/asio/spawn.hpp>
#include <boost/asio/steady_timer.hpp>
#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <chrono>

using namespace util;

struct CoroutineGroupTests : SyncAsioContextTest {
protected:
    testing::StrictMock<testing::MockFunction<void()>> callback1_;
    testing::StrictMock<testing::MockFunction<void()>> callback2_;
    testing::StrictMock<testing::MockFunction<void()>> callback3_;
};

TEST_F(CoroutineGroupTests, SpawnWait)
{
    testing::Sequence const sequence;
    EXPECT_CALL(callback1_, Call).InSequence(sequence);
    EXPECT_CALL(callback2_, Call).InSequence(sequence);
    EXPECT_CALL(callback3_, Call).InSequence(sequence);

    runSpawn([this](boost::asio::yield_context yield) {
        CoroutineGroup group{yield, 2};

        group.spawn(yield, [&](boost::asio::yield_context yield) {
            boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{1}};
            timer.async_wait(yield);
            callback1_.Call();
        });
        EXPECT_EQ(group.size(), 1);

        group.spawn(yield, [&](boost::asio::yield_context yield) {
            boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{2}};
            timer.async_wait(yield);
            callback2_.Call();
        });
        EXPECT_EQ(group.size(), 2);

        group.asyncWait(yield);
        EXPECT_EQ(group.size(), 0);

        callback3_.Call();
    });
}

TEST_F(CoroutineGroupTests, SpawnWaitSpawnWait)
{
    testing::Sequence const sequence;
    EXPECT_CALL(callback1_, Call).InSequence(sequence);
    EXPECT_CALL(callback2_, Call).InSequence(sequence);
    EXPECT_CALL(callback3_, Call).InSequence(sequence);

    runSpawn([this](boost::asio::yield_context yield) {
        CoroutineGroup group{yield, 2};

        group.spawn(yield, [&](boost::asio::yield_context yield) {
            boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{1}};
            timer.async_wait(yield);
            callback1_.Call();
        });
        EXPECT_EQ(group.size(), 1);

        group.asyncWait(yield);
        EXPECT_EQ(group.size(), 0);

        group.spawn(yield, [&](boost::asio::yield_context yield) {
            boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{1}};
            timer.async_wait(yield);
            callback2_.Call();
        });
        EXPECT_EQ(group.size(), 1);

        group.asyncWait(yield);
        EXPECT_EQ(group.size(), 0);

        callback3_.Call();
    });
}

TEST_F(CoroutineGroupTests, ChildCoroutinesFinishBeforeWait)
{
    testing::Sequence const sequence;
    EXPECT_CALL(callback2_, Call).InSequence(sequence);
    EXPECT_CALL(callback1_, Call).InSequence(sequence);
    EXPECT_CALL(callback3_, Call).InSequence(sequence);

    runSpawn([this](boost::asio::yield_context yield) {
        CoroutineGroup group{yield, 2};
        group.spawn(yield, [&](boost::asio::yield_context yield) {
            boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{2}};
            timer.async_wait(yield);
            callback1_.Call();
        });
        group.spawn(yield, [&](boost::asio::yield_context yield) {
            boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{1}};
            timer.async_wait(yield);
            callback2_.Call();
        });

        boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{3}};
        timer.async_wait(yield);

        group.asyncWait(yield);
        callback3_.Call();
    });
}

TEST_F(CoroutineGroupTests, EmptyGroup)
{
    EXPECT_CALL(callback1_, Call);

    runSpawn([this](boost::asio::yield_context yield) {
        CoroutineGroup group{yield};
        group.asyncWait(yield);
        callback1_.Call();
    });
}

TEST_F(CoroutineGroupTests, TooManyCoroutines)
{
    EXPECT_CALL(callback1_, Call);
    EXPECT_CALL(callback2_, Call);
    EXPECT_CALL(callback3_, Call);

    runSpawn([this](boost::asio::yield_context yield) {
        CoroutineGroup group{yield, 1};

        EXPECT_TRUE(group.spawn(yield, [this](boost::asio::yield_context innerYield) {
            boost::asio::steady_timer timer{
                innerYield.get_executor(), std::chrono::milliseconds{1}
            };
            timer.async_wait(innerYield);
            callback1_.Call();
        }));

        EXPECT_FALSE(group.spawn(yield, [this](boost::asio::yield_context) { callback2_.Call(); }));
        EXPECT_TRUE(group.isFull());

        boost::asio::steady_timer timer{yield.get_executor(), std::chrono::milliseconds{2}};
        timer.async_wait(yield);

        EXPECT_FALSE(group.isFull());
        EXPECT_TRUE(group.spawn(yield, [this](boost::asio::yield_context) { callback2_.Call(); }));

        group.asyncWait(yield);
        callback3_.Call();
    });
}

TEST_F(CoroutineGroupTests, SpawnForeign)
{
    testing::Sequence const sequence;
    EXPECT_CALL(callback1_, Call).InSequence(sequence);
    EXPECT_CALL(callback2_, Call).InSequence(sequence);

    runSpawn([this](boost::asio::yield_context yield) {
        CoroutineGroup group{yield, 1};

        auto const onForeignComplete = group.registerForeign(yield);
        [&]() { ASSERT_TRUE(onForeignComplete.has_value()); }();

        [&]() { ASSERT_FALSE(group.registerForeign(yield).has_value()); }();

        util::spawn(ctx_, [this, &onForeignComplete](boost::asio::yield_context innerYield) {
            boost::asio::steady_timer timer{
                innerYield.get_executor(), std::chrono::milliseconds{2}
            };
            timer.async_wait(innerYield);
            callback1_.Call();
            onForeignComplete->operator()();
        });

        group.asyncWait(yield);
        callback2_.Call();
    });
}