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

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

using namespace NSrvKernel;

namespace {

constexpr size_t ITER_COUNT = 5;

}  // namespace

Y_UNIT_TEST_SUITE(TCoroSingleCondVarTest) {
    CORO_TEST_BEGIN(SmokingTest, exec) {
        bool started = false;
        TCoroSingleCondVar cv;

        // Does nothing.
        cv.notify();

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

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

        cv.notify();
        UNIT_ASSERT_EQUAL(fut.Get(), true);
    } CORO_TEST_END

    CORO_TEST_BEGIN(MultipassTest, exec) {
        TCoroSingleCondVar cv;
        size_t iterNum = 0;

        TCoroutine task{"MultipassTest", exec, [&] {
            while (!exec->Running()->Cancelled()) {
                ++iterNum;
                Y_UNUSED(cv.wait(exec));
            }
        }};

        for (size_t i = 0; i < ITER_COUNT; ++i) {
            exec->Running()->Yield();
            cv.notify();
        }
        task.Cancel();
        task.Join();
        UNIT_ASSERT_EQUAL(iterNum, ITER_COUNT);
    } CORO_TEST_END
}

Y_UNIT_TEST_SUITE(TCoroCondVarTest) {
    CORO_TEST_BEGIN(ManyWaitingCoroutinesTest, exec) {
        size_t result = 0;
        TVector<TCoroutine> tasks;
        TCoroCondVar cv;
        auto isRunning = std::mem_fn(&TCoroutine::Running);

        // Does nothing.
        cv.notify_one();
        cv.notify_all();

        std::generate_n(std::back_inserter(tasks), ITER_COUNT, [&cv, &result, exec] {
            return TCoroutine{"ManyWaitingCoroutinesTest", exec, [&cv, &result, exec] {
                Y_UNUSED(cv.wait(exec));
                ++result;
            }};
        });
        UNIT_ASSERT(std::all_of(tasks.begin(), tasks.end(), isRunning));

        UNIT_ASSERT_EQUAL(result, 0);
        // Make all coroutines wait for condvar.
        exec->Running()->Yield();
        UNIT_ASSERT_EQUAL(result, 0);

        cv.notify_one();
        exec->Running()->Yield();
        UNIT_ASSERT_EQUAL(result, 1);
        UNIT_ASSERT(!tasks.front().Running());
        UNIT_ASSERT(std::all_of(tasks.begin() + 1, tasks.end(), isRunning));

        cv.notify_all();
        exec->Running()->Yield();
        UNIT_ASSERT_EQUAL(result, ITER_COUNT);
        UNIT_ASSERT(std::none_of(tasks.begin(), tasks.end(), isRunning));
    } CORO_TEST_END

    CORO_TEST_BEGIN(WaitingCoroutinesCanceledTest, exec) {
        TVector<TCoroFuture<int>> results;
        TCoroCondVar cv;
        auto isValid = std::mem_fn(&TCoroFuture<int>::Valid);

        std::generate_n(std::back_inserter(results), ITER_COUNT, [&cv, exec] {
            return CoroAsync("WaitingCoroutinesCanceledTest", exec, [&cv, exec] {
                return cv.wait(exec);
            });
        });
        UNIT_ASSERT(std::all_of(results.begin(), results.end(), isValid));

        // Make all coroutines wait for condvar.
        exec->Running()->Yield();
        for (auto& f : results) {
            f.Cancel();
        }
        exec->Running()->Yield();
        UNIT_ASSERT(std::all_of(results.begin(), results.end(), isValid));
        UNIT_ASSERT(std::all_of(results.begin(), results.end(), [](auto& f) {
            return f.Get() == ECANCELED;
        }));
        UNIT_ASSERT(std::none_of(results.begin(), results.end(), isValid));
    } CORO_TEST_END
}
