#include "test_utils.hpp"

#include <ydb-cpp-sdk/library/issue/yql_issue.h>

#include <ydb/impl/future.hpp>
#include <ydb/impl/retry.hpp>

USERVER_NAMESPACE_BEGIN

namespace {

class RetryOperationFixture : public ydb::ClientFixtureBase {
public:
    template <typename Func>
    auto RetryOperationSync(std::size_t retries, Func func) {
        auto settings = MakeOperationSettings(retries);
        ydb::impl::RequestContext context{GetTableClient(), ydb::Query{}, settings};

        auto future = ydb::impl::RetryOperation(context, std::move(func));

        return ydb::impl::GetFutureValueUnchecked(std::move(future));
    }

private:
    ydb::OperationSettings MakeOperationSettings(std::uint32_t retries) {
        static constexpr std::chrono::milliseconds kTimeout{3000};
        return ydb::OperationSettings{
            .retries = retries,
            .operation_timeout_ms = kTimeout,
            .cancel_after_ms = kTimeout,
            .client_timeout_ms = kTimeout,
            .get_session_timeout_ms = kTimeout,
        };
    }
};

constexpr NYdb::EStatus kSuccess = NYdb::EStatus::SUCCESS;
constexpr NYdb::EStatus kRetryableStatus = NYdb::EStatus::ABORTED;
constexpr NYdb::EStatus kNonRetryableStatus = NYdb::EStatus::BAD_REQUEST;

inline NThreading::TFuture<NYdb::TStatus> MakeStatusFuture(NYdb::EStatus status) {
    return NThreading::MakeFuture<NYdb::TStatus>(NYdb::TStatus{status, NYdb::NIssue::TIssues{}});
}

class TestOperationResults final : public NYdb::TStatus {
public:
    TestOperationResults(NYdb::TStatus&& status, const std::string& data)
        : NYdb::TStatus(std::move(status)), data_(data) {}

    TestOperationResults(TestOperationResults&&) = default;

    const std::string& GetData() const { return data_; }

private:
    std::string data_;
};

}  // namespace

UTEST_F(RetryOperationFixture, HandleOfInheritorsOfTStatus) {
    const std::string data = "qwerty";
    const auto res = RetryOperationSync(
        /*retries=*/0,
        [&data](NYdb::NTable::TSession) {
            return NThreading::MakeFuture<TestOperationResults>(TestOperationResults{
                NYdb::TStatus{kSuccess, NYdb::NIssue::TIssues{}}, data});
        }
    );
    ASSERT_EQ(res.GetData(), data);
};

UTEST_F(RetryOperationFixture, Success) {
    std::size_t attempts = 0;
    const auto res = RetryOperationSync(/*retries=*/3, [&attempts](NYdb::NTable::TSession) {
        attempts++;
        return MakeStatusFuture(kSuccess);
    });

    ASSERT_EQ(res.GetStatus(), NYdb::EStatus::SUCCESS);
    ASSERT_EQ(attempts, 1);
};

UTEST_F(RetryOperationFixture, NonRetry) {
    std::size_t attempts = 0;
    const auto res = RetryOperationSync(/*retries=*/3, [&attempts](NYdb::NTable::TSession) {
        attempts++;
        return MakeStatusFuture(kNonRetryableStatus);
    });
    ASSERT_EQ(res.GetStatus(), kNonRetryableStatus);
    ASSERT_EQ(attempts, 1);
};

UTEST_F(RetryOperationFixture, SuccessOnTheLastAttempt) {
    constexpr std::uint32_t kRetries = 5;
    std::size_t attempts = 0;
    const auto res = RetryOperationSync(/*retries=*/kRetries, [&attempts](NYdb::NTable::TSession) {
        attempts++;
        if (attempts < kRetries) {
            return MakeStatusFuture(kRetryableStatus);
        }
        return MakeStatusFuture(kSuccess);
    });
    ASSERT_EQ(res.GetStatus(), NYdb::EStatus::SUCCESS);
    ASSERT_EQ(attempts, kRetries);
};

UTEST_F(RetryOperationFixture, AttemptsIsRetriesPlusOne) {
    constexpr std::uint32_t kRetries = 5;
    std::size_t attempts = 0;
    const auto res = RetryOperationSync(
        /*retries=*/kRetries,
        [&attempts](NYdb::NTable::TSession) {
            attempts++;
            return MakeStatusFuture(kRetryableStatus);
        }
    );
    ASSERT_EQ(attempts, kRetries + 1);
    ASSERT_EQ(res.GetStatus(), kRetryableStatus);
};

UTEST_F(RetryOperationFixture, RetriesLimit) {
    // ydb-sdk has own maximum for retries, so we want to step over this
    constexpr std::uint32_t kRetries = 1000;
    std::size_t attempts = 0;
    const auto res = RetryOperationSync(
        /*retries=*/1000,
        [&attempts](NYdb::NTable::TSession) {
            attempts++;
            if (attempts < kRetries) {
                return MakeStatusFuture(kRetryableStatus);
            }
            // Unreachable!
            return MakeStatusFuture(kSuccess);
        }
    );
    ASSERT_EQ(res.GetStatus(), kRetryableStatus);
};

UTEST_F(RetryOperationFixture, Exception) {
    UASSERT_THROW_MSG(
        RetryOperationSync(
            /*retries=*/0,
            [](NYdb::NTable::TSession) {
                throw std::runtime_error{"error"};
                return MakeStatusFuture(kSuccess);
            }
        ),
        std::runtime_error,
        "error"
    );
};

USERVER_NAMESPACE_END
