#include <balancer/kernel/coro/coro_event.h>
#include <balancer/kernel/coro/coro_async.h>
#include <balancer/kernel/coro/coroutine.h>

#include <balancer/kernel/testing/testing.h>

#include <numeric>

using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TCoroManualEventTest) {
    constexpr size_t ITER_COUNT = 5;

    CORO_TEST_BEGIN(NotifyBeforeTest, exec) {
        bool done = false;
        TCoroManualEvent event;

        event.notify();

        auto fut = CoroAsync("NotifyBeforeTest", exec, [&] {
            UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
            return done = true;
        });

        UNIT_ASSERT(!done);
        UNIT_ASSERT_EQUAL(fut.Get(), true);
        UNIT_ASSERT(done);
    } CORO_TEST_END

    CORO_TEST_BEGIN(NotifyAfterTest, exec) {
        bool started = false;
        bool done = false;
        TCoroManualEvent event;

        auto fut = CoroAsync("NotifyAfterTest", exec, [&] {
            started = true;
            UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
            return done = true;
        });
        UNIT_ASSERT(!started);

        exec->Running()->Yield();
        UNIT_ASSERT(started);
        UNIT_ASSERT(!done);

        event.notify();
        UNIT_ASSERT_EQUAL(fut.Get(), true);
        UNIT_ASSERT(done);
    } CORO_TEST_END

    CORO_TEST_BEGIN(MultipassTest, exec) {
        TCoroManualEvent event;
        bool started = false;

        auto fut = CoroAsync("MultipassTest", exec, [&] {
            started = true;
            TVector<int> ret;
            for (size_t i = 0; i < ITER_COUNT; ++i) {
                UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
                ret.push_back(i);
            }
            return ret;
        });
        UNIT_ASSERT(!started);

        exec->Running()->Yield();
        UNIT_ASSERT(started);

        event.notify();
        TVector<int> etalon(ITER_COUNT);
        std::iota(etalon.begin(), etalon.end(), 0);
        UNIT_ASSERT_EQUAL(fut.Get(), etalon);
    } CORO_TEST_END

    CORO_TEST_BEGIN(ResetTest, exec) {
        TCoroManualEvent event;

        TCoroutine task{"ResetTest", exec, &TCoroManualEvent::notify, &event};
        UNIT_ASSERT(task.Running());

        UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
        UNIT_ASSERT(!task.Running());

        for (size_t i = 0; i < ITER_COUNT; ++i) {
            // No wait actually.
            UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
        }

        task = TCoroutine{"ResetTest", exec, &TCoroManualEvent::notify, &event};
        UNIT_ASSERT(task.Running());

        event.Reset();

        UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
        UNIT_ASSERT(!task.Running());
    } CORO_TEST_END
}

Y_UNIT_TEST_SUITE(TCoroAutoEventTest) {
    constexpr size_t ITER_COUNT = 5;

    CORO_TEST_BEGIN(NotifyBeforeTest, exec) {
        bool done = false;
        TCoroAutoEvent event;

        event.notify();

        auto fut = CoroAsync("NotifyBeforeTest", exec, [&] {
            UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
            return done = true;
        });

        UNIT_ASSERT(!done);
        UNIT_ASSERT_EQUAL(fut.Get(), true);
        UNIT_ASSERT(done);
    } CORO_TEST_END

    CORO_TEST_BEGIN(NotifyAfterTest, exec) {
        bool started = false;
        bool done = false;
        TCoroAutoEvent event;

        auto fut = CoroAsync("NotifyAfterTest", exec, [&] {
            started = true;
            UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
            return done = true;
        });
        UNIT_ASSERT(!started);

        exec->Running()->Yield();
        UNIT_ASSERT(started);
        UNIT_ASSERT(!done);

        event.notify();
        UNIT_ASSERT_EQUAL(fut.Get(), true);
        UNIT_ASSERT(done);
    } CORO_TEST_END

    CORO_TEST_BEGIN(MultipassTest, exec) {
        TCoroAutoEvent event;
        bool started = false;

        auto fut = CoroAsync("MultipassTest", exec, [&] {
            started = true;
            TVector<int> ret;
            for (size_t i = 0; i < ITER_COUNT; ++i) {
                UNIT_ASSERT_EQUAL(event.wait(exec), EWAKEDUP);
                ret.push_back(i);
            }
            return ret;
        });
        UNIT_ASSERT(!started);

        exec->Running()->Yield();
        UNIT_ASSERT(started);

        for (size_t i = 0; i < ITER_COUNT; ++i) {
            event.notify();
            exec->Running()->Yield();
        }
        event.notify();
        TVector<int> etalon(ITER_COUNT);
        std::iota(etalon.begin(), etalon.end(), 0);
        UNIT_ASSERT_EQUAL(fut.Get(), etalon);
    } CORO_TEST_END
}
