#include "data/cassandra/Error.hpp"
#include "data/cassandra/FakesAndMocks.hpp"
#include "data/cassandra/impl/AsyncExecutor.hpp"
#include "util/AsioContextTestFixture.hpp"

#include <boost/asio/executor_work_guard.hpp>
#include <boost/asio/io_context.hpp>
#include <cassandra.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include <atomic>
#include <functional>
#include <optional>
#include <thread>
#include <utility>

using namespace data::cassandra;
using namespace data::cassandra::impl;
using namespace testing;

class BackendCassandraAsyncExecutorTest : public SyncAsioContextTest {
protected:
    struct CallbackMock {
        MOCK_METHOD(void, onComplete, (FakeResultOrError));
        MOCK_METHOD(void, onRetry, ());
    };
    CallbackMock callbackMock_;
    std::function<void()> onRetry_ = [this]() { callbackMock_.onRetry(); };
};

TEST_F(BackendCassandraAsyncExecutorTest, CompletionCalledOnSuccess)
{
    auto handle = MockHandle{};

    ON_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .WillByDefault([this](auto const&, auto&& cb) {
            boost::asio::post(ctx_, [cb = std::forward<decltype(cb)>(cb)]() { cb({}); });
            return FakeFutureWithCallback{};
        });
    EXPECT_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .Times(AtLeast(1));

    auto work = std::make_optional(boost::asio::make_work_guard(ctx_));
    EXPECT_CALL(callbackMock_, onComplete);

    AsyncExecutor<FakeStatement, MockHandle>::run(
        ctx_,
        handle,
        FakeStatement{},
        [&work, this](auto resultOrError) {
            callbackMock_.onComplete(std::move(resultOrError));
            work.reset();
        },
        std::move(onRetry_)
    );

    ctx_.run();
}

TEST_F(BackendCassandraAsyncExecutorTest, ExecutedMultipleTimesByRetryPolicyOnMainThread)
{
    auto callCount = std::atomic_int{0};
    auto handle = MockHandle{};

    // emulate successful execution after some attempts
    ON_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .WillByDefault([&callCount](auto const&, auto&& cb) {
            ++callCount;
            if (callCount >= 3) {
                cb({});
            } else {
                cb({CassandraError{"timeout", CASS_ERROR_LIB_REQUEST_TIMED_OUT}});
            }

            return FakeFutureWithCallback{};
        });
    EXPECT_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .Times(3);

    auto work = std::make_optional(boost::asio::make_work_guard(ctx_));
    EXPECT_CALL(callbackMock_, onComplete);
    EXPECT_CALL(callbackMock_, onRetry).Times(2);

    AsyncExecutor<FakeStatement, MockHandle>::run(
        ctx_,
        handle,
        FakeStatement{},
        [this, &work](auto resultOrError) {
            callbackMock_.onComplete(std::move(resultOrError));
            work.reset();
        },
        std::move(onRetry_)
    );

    ctx_.run();
    ASSERT_EQ(callCount, 3);
}

TEST_F(BackendCassandraAsyncExecutorTest, ExecutedMultipleTimesByRetryPolicyOnOtherThread)
{
    auto callCount = std::atomic_int{0};
    auto handle = MockHandle{};

    auto threadedCtx = boost::asio::io_context{};
    auto work = std::make_optional(boost::asio::make_work_guard(threadedCtx));
    auto thread = std::thread{[&threadedCtx] { threadedCtx.run(); }};

    // emulate successful execution after some attempts
    ON_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .WillByDefault([&callCount](auto const&, auto&& cb) {
            ++callCount;
            if (callCount >= 3) {
                cb({});
            } else {
                cb({CassandraError{"timeout", CASS_ERROR_LIB_REQUEST_TIMED_OUT}});
            }

            return FakeFutureWithCallback{};
        });
    EXPECT_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .Times(3);

    auto work2 = std::make_optional(boost::asio::make_work_guard(ctx_));
    EXPECT_CALL(callbackMock_, onComplete);
    EXPECT_CALL(callbackMock_, onRetry).Times(2);

    AsyncExecutor<FakeStatement, MockHandle>::run(
        threadedCtx,
        handle,
        FakeStatement{},
        [this, &work, &work2](auto resultOrError) {
            callbackMock_.onComplete(std::move(resultOrError));
            work.reset();
            work2.reset();
        },
        std::move(onRetry_)
    );

    ctx_.run();
    EXPECT_EQ(callCount, 3);
    threadedCtx.stop();
    thread.join();
}

TEST_F(BackendCassandraAsyncExecutorTest, CompletionCalledOnFailureAfterRetryCountExceeded)
{
    auto handle = MockHandle{};

    // FakeRetryPolicy returns false for shouldRetry in which case we should
    // still call onComplete giving it whatever error we have raised internally.
    ON_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .WillByDefault([](auto const&, auto&& cb) {
            cb({CassandraError{"not a timeout", CASS_ERROR_LIB_INTERNAL_ERROR}});
            return FakeFutureWithCallback{};
        });
    EXPECT_CALL(
        handle,
        asyncExecute(An<FakeStatement const&>(), An<std::function<void(FakeResultOrError)>&&>())
    )
        .Times(1);

    auto work = std::make_optional(boost::asio::make_work_guard(ctx_));
    EXPECT_CALL(callbackMock_, onComplete);

    AsyncExecutor<FakeStatement, MockHandle, FakeRetryPolicy>::run(
        ctx_,
        handle,
        FakeStatement{},
        [this, &work](auto res) {
            EXPECT_FALSE(res);
            EXPECT_EQ(res.error().code(), CASS_ERROR_LIB_INTERNAL_ERROR);
            EXPECT_EQ(res.error().message(), "not a timeout");

            callbackMock_.onComplete(std::move(res));
            work.reset();
        },
        std::move(onRetry_)
    );

    ctx_.run();
}