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

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

using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TCoroutineTest) {
    constexpr size_t ITER_COUNT = 5;

    CORO_TEST_BEGIN(SmokingTest, exec) {
        bool flag = false;
        TCoroutine task("SmokingTest", exec, [&]() {
            flag = true;
        });
        UNIT_ASSERT(task.Running());
        exec->Running()->Yield();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(flag);
    } CORO_TEST_END

    CORO_TEST_BEGIN(WithArgsTest, exec) {
        int res = 0;
        TCoroutine task("WithArgsTest", exec, [&](int a, int b, int c) {
            res = a + b + c;
        }, 5, 6, 7);
        UNIT_ASSERT(task.Running());
        exec->Running()->Yield();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT_EQUAL(res, 5 + 6 + 7);
    } CORO_TEST_END

    struct TSomeStruct {
        void Foo() {
            Flag = true;
        }

        bool Flag = false;
    };

    CORO_TEST_BEGIN(ClassMethodTest, exec) {
        TSomeStruct obj;
        TCoroutine task("ClassMethodTest", exec, &TSomeStruct::Foo, &obj);
        UNIT_ASSERT(task.Running());
        exec->Running()->Yield();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(obj.Flag);
    } CORO_TEST_END

    static void SomeFunc(bool* flag) {
        *flag = true;
    }

    CORO_TEST_BEGIN(StaticFunctionTest, exec) {
        bool flag = false;
        TCoroutine task("StaticFunctionTest", exec, SomeFunc, &flag);
        UNIT_ASSERT(task.Running());
        exec->Running()->Yield();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(flag);
    } CORO_TEST_END

    CORO_TEST_BEGIN(ChainTest, exec) {
        TVector<int> v;
        TCoroutine task1("ChainTest1", exec, [](TVector<int>* v, TContExecutor* const exec) {
            TCoroutine task2("ChainTest2", exec, [](TVector<int>* v, TContExecutor* const exec) {
                TCoroutine task3("ChainTest3", exec, [](TVector<int>* v) {
                    v->push_back(3);
                }, v);
                task3.Join();
                v->push_back(2);
            }, v, exec);
            task2.Join();
            v->push_back(1);
        }, &v, exec);
        UNIT_ASSERT(task1.Running());
        task1.Join();
        UNIT_ASSERT(!task1.Running());
        v.push_back(0);
        UNIT_ASSERT(std::is_sorted(v.rbegin(), v.rend()));
    } CORO_TEST_END

    CORO_TEST_BEGIN(SwitchingTest, exec) {
        TVector<int> v;
        TCoroutine task("SwitchingTest1", exec, [&] {
            TCoroutine task1("SwitchingTest2", exec, [&] {
                for (size_t i = 0; i < ITER_COUNT; ++i) {
                    v.push_back(2);
                    exec->Running()->Yield();
                }
            });
            for (size_t i = 0; i < ITER_COUNT; ++i) {
                v.push_back(1);
                exec->Running()->Yield();
            }
        });
        UNIT_ASSERT(task.Running());
        task.Join();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT_EQUAL(std::adjacent_find(v.begin(), v.end()), v.end());
    } CORO_TEST_END

    CORO_TEST_BEGIN(CancelTest, exec) {
        TCoroutine task("CancelTest", exec, [&] {
            auto* const cont = exec->Running();
            while (!cont->Cancelled()) {
                cont->Yield();
            }
        });
        for (size_t i = 0; i < ITER_COUNT; ++i) {
            UNIT_ASSERT(task.Running());
            exec->Running()->Yield();
        }
        task.Cancel();
        UNIT_ASSERT(task.Running());
        task.Join();
        UNIT_ASSERT(!task.Running());
    } CORO_TEST_END

    CORO_TEST_BEGIN(ResetTest, exec) {
        bool flag = false;
        TCoroutine task("ResetTest", exec, [&] {
            auto* const cont = exec->Running();
            while (!cont->Cancelled()) {
                cont->Yield();
            }
            flag = true;
        });
        for (size_t i = 0; i < ITER_COUNT; ++i) {
            UNIT_ASSERT(task.Running());
            UNIT_ASSERT(!flag);
            exec->Running()->Yield();
        }
        task = {};
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(flag);
    } CORO_TEST_END

    CORO_TEST_BEGIN(MoveTest, exec) {
        bool flag = false;
        TCoroutine task("MoveTest1", exec, [&] {
            flag = true;
        });
        UNIT_ASSERT(task.Running());
        UNIT_ASSERT(!flag);

        // Move construction test
        // --------------------------------------------------------------------------------
        TCoroutine task2 = std::move(task);
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(task2.Running());
        UNIT_ASSERT(!flag);

        task2.Join();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(!task2.Running());
        UNIT_ASSERT(flag);
        // --------------------------------------------------------------------------------

        task = TCoroutine{"MoveTest2", exec, [&] {
            flag = false;
        }};
        UNIT_ASSERT(task.Running());
        UNIT_ASSERT(!task2.Running());
        UNIT_ASSERT(flag);

        // Move assignment test
        // --------------------------------------------------------------------------------
        task2 = std::move(task);
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(task2.Running());
        UNIT_ASSERT(flag);

        exec->Running()->Yield();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(!task2.Running());
        UNIT_ASSERT(!flag);
        // --------------------------------------------------------------------------------
    } CORO_TEST_END

    struct TWidget : public TMoveOnly {};

    CORO_TEST_BEGIN(MoveArgTest, exec) {
        bool flag = false;
        TWidget w;
        TCoroutine task("MoveArgTest", exec, [&](TWidget&&) {
            flag = true;
        }, std::move(w));
        UNIT_ASSERT(task.Running());
        UNIT_ASSERT(!flag);

        task.Join();
        UNIT_ASSERT(!task.Running());
        UNIT_ASSERT(flag);
    } CORO_TEST_END

    Y_UNIT_TEST(Exception) {
        TContExecutor exec(32000);

        TCoroutine coroutine("throw", &exec, []() {
            ythrow yexception();
        });

        exec.Execute();
    }
}
