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

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

using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TCoroAsyncTest) {
    constexpr size_t ITER_COUNT = 5;
    constexpr int A_VAL_SAMPLE = 23;
    constexpr int B_VAL_SAMPLE = 58;
    constexpr int C_VAL_SAMPLE = 93;

    Y_UNIT_TEST(StaticTest) {
        using TFuture1 = decltype(CoroAsync("", nullptr, std::declval<int()>()));
        static_assert(std::is_same_v<TFuture1, TCoroFuture<int>>);
    }

    CORO_TEST_BEGIN(SmokingTest, exec) {
        auto fut = CoroAsync("SmokingTest", exec, [] {
            return true;
        });
        UNIT_ASSERT(fut.Valid());
        UNIT_ASSERT(fut.Get());
        UNIT_ASSERT(!fut.Valid());
    } CORO_TEST_END

    Y_UNIT_TEST(DefaultTest) {
        TCoroFuture<int> fut;
        UNIT_ASSERT(!fut.Valid());
    }

    CORO_TEST_BEGIN(WithArgsTest, exec) {
        bool started = false;
        auto fut = CoroAsync("WithArgsTest", exec, [&](int a, int b, int c) {
            started = true;
            return a + b + c;
        }, A_VAL_SAMPLE, B_VAL_SAMPLE, C_VAL_SAMPLE);
        UNIT_ASSERT(fut.Valid());
        UNIT_ASSERT(!started);

        fut.Wait();
        UNIT_ASSERT(fut.Valid());
        UNIT_ASSERT(started);

        UNIT_ASSERT_EQUAL(fut.Get(), A_VAL_SAMPLE + B_VAL_SAMPLE + C_VAL_SAMPLE);
        UNIT_ASSERT(!fut.Valid());
        UNIT_ASSERT(started);
    } CORO_TEST_END

    struct TSomeStruct {
        bool Foo() {
            return true;
        }
    };

    CORO_TEST_BEGIN(ClassMethodTest, exec) {
        TSomeStruct obj;
        auto fut = CoroAsync("ClassMethodTest", exec, &TSomeStruct::Foo, &obj);
        UNIT_ASSERT(fut.Valid());

        exec->Running()->Yield();
        UNIT_ASSERT(fut.Valid());

        UNIT_ASSERT(fut.Get());
        UNIT_ASSERT(!fut.Valid());
    } CORO_TEST_END

    static int SomeFunc(int a, int b) {
        return a + b;
    }

    CORO_TEST_BEGIN(StaticFunctionTest, exec) {
        auto fut = CoroAsync("StaticFunctionTest", exec,
                             SomeFunc, A_VAL_SAMPLE, B_VAL_SAMPLE);
        UNIT_ASSERT(fut.Valid());
        UNIT_ASSERT_EQUAL(fut.Get(), A_VAL_SAMPLE + B_VAL_SAMPLE);
        UNIT_ASSERT(!fut.Valid());
    } CORO_TEST_END

    CORO_TEST_BEGIN(ChainTest, exec) {
        TVector<int> v;
        auto fut1 = CoroAsync("ChainTest1", exec, [](TContExecutor* const e) {
            auto fut2 = CoroAsync("ChainTest2", e, [](TContExecutor* const e) {
                auto fut3 = CoroAsync("ChainTest3", e, [] {
                    return TVector<int>{3};
                });
                UNIT_ASSERT(fut3.Valid());

                auto ret = fut3.Get();
                UNIT_ASSERT(!fut3.Valid());

                ret.push_back(2);
                return ret;
            }, e);
            UNIT_ASSERT(fut2.Valid());

            fut2.Wait();
            UNIT_ASSERT(fut2.Valid());

            auto ret = fut2.Get();
            UNIT_ASSERT(!fut2.Valid());

            ret.push_back(1);
            return ret;
        }, exec);
        UNIT_ASSERT(fut1.Valid());

        fut1.Wait();
        UNIT_ASSERT(fut1.Valid());

        auto ret = fut1.Get();
        v.push_back(0);
        UNIT_ASSERT(std::is_sorted(v.rbegin(), v.rend()));
    } CORO_TEST_END

    CORO_TEST_BEGIN(SwitchingTest, exec) {
        TVector<int> v;
        auto fut = CoroAsync("SwitchingTest1", exec, [&] {
            auto fut1 = CoroAsync("SwitchingTest2", exec, [&] {
                for (size_t i = 0; i < ITER_COUNT; ++i) {
                    v.push_back(2);
                    exec->Running()->Yield();
                }
                return true;
            });
            UNIT_ASSERT(fut1.Valid());
            for (size_t i = 0; i < ITER_COUNT; ++i) {
                v.push_back(1);
                exec->Running()->Yield();
            }
            return fut1.Get();
        });
        UNIT_ASSERT(fut.Valid());
        fut.Wait();

        UNIT_ASSERT(fut.Valid());
        UNIT_ASSERT_EQUAL(std::adjacent_find(v.begin(), v.end()), v.end());

        UNIT_ASSERT(fut.Get());
        UNIT_ASSERT(!fut.Valid());
    } CORO_TEST_END

    CORO_TEST_BEGIN(CancelTest, exec) {
        auto fut = CoroAsync("CancelTest", exec, [](TContExecutor* const e) {
            auto* cont = e->Running();
            int ret = 0;
            while (!cont->Cancelled()) {
                cont->Yield();
                ++ret;
            }
            return ret;
        }, exec);
        for (size_t i = 0; i < ITER_COUNT; ++i) {
            UNIT_ASSERT(fut.Valid());
            exec->Running()->Yield();
        }
        fut.Cancel();
        UNIT_ASSERT(fut.Valid());

        UNIT_ASSERT_EQUAL(fut.Get(), ITER_COUNT);
        UNIT_ASSERT(!fut.Valid());
    } CORO_TEST_END

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

    CORO_TEST_BEGIN(MoveTest, exec) {
        bool flag = false;
        auto fut = CoroAsync("MoveTest1", exec, [&] {
            return flag = true;
        });
        UNIT_ASSERT(fut.Valid());
        UNIT_ASSERT(!flag);

        // Move construction test
        // --------------------------------------------------------------------------------
        auto fut2 = std::move(fut);
        UNIT_ASSERT(!fut.Valid());
        UNIT_ASSERT(fut2.Valid());
        UNIT_ASSERT(!flag);

        UNIT_ASSERT(fut2.Get());
        UNIT_ASSERT(!fut.Valid());
        UNIT_ASSERT(!fut2.Valid());
        UNIT_ASSERT(flag);
        // --------------------------------------------------------------------------------

        fut = CoroAsync("MoveTest2", exec, [&] {
            return flag = false;
        });
        UNIT_ASSERT(fut.Valid());
        UNIT_ASSERT(!fut2.Valid());
        UNIT_ASSERT(flag);

        // Move assignment test
        // --------------------------------------------------------------------------------
        fut2 = std::move(fut);
        UNIT_ASSERT(!fut.Valid());
        UNIT_ASSERT(fut2.Valid());
        UNIT_ASSERT(flag);

        exec->Running()->Yield();
        UNIT_ASSERT(!fut.Valid());
        UNIT_ASSERT(fut2.Valid());
        UNIT_ASSERT(!flag);

        UNIT_ASSERT(!fut2.Get());
        UNIT_ASSERT(!fut.Valid());
        UNIT_ASSERT(!fut2.Valid());
        UNIT_ASSERT(!flag);
        // --------------------------------------------------------------------------------
    } CORO_TEST_END

    CORO_TEST_BEGIN(SwapTest, exec) {
        auto fut1 = CoroAsync("SwapTest1", exec, [] {
            return A_VAL_SAMPLE;
        });
        UNIT_ASSERT(fut1.Valid());
        auto fut2 = CoroAsync("SwapTest1", exec, [] {
            return B_VAL_SAMPLE;
        });
        UNIT_ASSERT(fut2.Valid());

        fut1.Swap(fut2);
        UNIT_ASSERT(fut1.Valid());
        UNIT_ASSERT(fut2.Valid());

        UNIT_ASSERT_EQUAL(fut1.Get(), B_VAL_SAMPLE);
        UNIT_ASSERT_EQUAL(fut2.Get(), A_VAL_SAMPLE);
        UNIT_ASSERT(!fut1.Valid());
        UNIT_ASSERT(!fut2.Valid());
    } CORO_TEST_END

    struct TWidget : public TMoveOnly {};

    CORO_TEST_BEGIN(MoveArgTest, exec) {
        TWidget w;
        auto fut = CoroAsync("MoveArgTest", exec, [&](TWidget w) {
            return w;
        }, std::move(w));
        UNIT_ASSERT(fut.Valid());

        w = fut.Get();
        UNIT_ASSERT(!fut.Valid());
    } CORO_TEST_END

    CORO_TEST_BEGIN(DestructionTest, exec) {
        bool done = false;
        {
            auto fut = CoroAsync("DestructionTest", exec, [&] {
                UNIT_ASSERT_EQUAL(exec->Running()->SleepI(), ECANCELED);
                return done = true;
            });
            UNIT_ASSERT(fut.Valid());
        }
        UNIT_ASSERT(done);
    } CORO_TEST_END

    CORO_TEST_BEGIN(ThrowTest, exec) {
        auto fut = CoroAsync("ThrowTest", exec, []() -> int {
            throw TSystemError{};
        });

        UNIT_ASSERT_EXCEPTION(fut.Get(), TSystemError);
    } CORO_TEST_END
}
