#include <balancer/kernel/balancer/ut_mocks/algorithm_mock.h>
#include <balancer/kernel/balancer/ut_mocks/backend_mock.h>
#include <balancer/kernel/balancer/ut_mocks/policy_mock.h>

#include <balancer/modules/balancer/policies/simple.h>
#include <balancer/modules/balancer/policies/unique.h>
#include <balancer/modules/balancer/policies/retry.h>
#include <balancer/modules/balancer/policies/unique_retry.h>
#include <balancer/modules/balancer/policies/timeout.h>
#include <balancer/modules/balancer/policies/active.h>

#include <library/cpp/testing/unittest/registar.h>

#include <util/generic/xrange.h>

using namespace NSrvKernel;

using ::testing::Return;
using ::testing::_;
using ::testing::Lt;

Y_UNIT_TEST_SUITE(TPolicyUnittest) {
    constexpr size_t ITERATIONS_NUM = 10;

    Y_UNIT_TEST(SimplePolicyTest) {
        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        NSimple::TSimplePolicy policy;


        // On Next call policy just returns result of SelectNext in algorithm
        // --------------------------------------------------------------------------------
        EXPECT_CALL(algoMock, SelectNext())
            .Times(2)
            .WillOnce(Return(nullptr))
            .WillOnce(Return(&backMock));

        auto* backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, nullptr);

        backend = policy.Next(&algoMock, true);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(UniquePolicyTest) {
        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        NUnique::TUniquePolicy policy;


        // Policy calls RemoveSelected when algorithm finds backend
        // --------------------------------------------------------------------------------
        EXPECT_CALL(algoMock, Next())
            .Times(1)
            .WillOnce(Return(&backMock));

        EXPECT_CALL(algoMock, RemoveSelected(&backMock))
            .Times(1);

        auto* backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------


        // Policy doesn't call RemoveSelected when algorithm can't find any backend
        // --------------------------------------------------------------------------------
        EXPECT_CALL(algoMock, Next())
            .Times(1)
            .WillOnce(Return(nullptr));

        EXPECT_CALL(algoMock, RemoveSelected(&backMock))
            .Times(0);

        backend = policy.Next(&algoMock, true);
        UNIT_ASSERT_EQUAL(backend, nullptr);
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(RetryPolicyTest) {
        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        auto policyPtr = MakeHolder<NSrvKernel::NTesting::TPolicyMock>();
        auto* policyMock = policyPtr.Get();

        NRetry::TRetryPolicy policy{ std::move(policyPtr) };


        // Situation when algorithm finds backend instantly
        // --------------------------------------------------------------------------------
        EXPECT_CALL(*policyMock, Next(&algoMock, false))
            .Times(1)
            .WillOnce(Return(&backMock));

        EXPECT_CALL(algoMock, Reset())
            .Times(0);

        auto* backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------


        // Situation when algorithm finds any backend only after second call
        // So algorithm will be reseted after first failure
        // --------------------------------------------------------------------------------
        EXPECT_CALL(*policyMock, Next(&algoMock, true))
            .Times(2)
            .WillOnce(Return(nullptr))
            .WillOnce(Return(&backMock));

        EXPECT_CALL(algoMock, Reset())
            .Times(1);

        backend = policy.Next(&algoMock, true);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------


        // Basic invariants
        // --------------------------------------------------------------------------------
        EXPECT_CALL(*policyMock, RegisterSuccess())
            .Times(1);

        policy.RegisterSuccess();

        EXPECT_CALL(*policyMock, RegisterFail())
            .Times(1);

        policy.RegisterFail();
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(RetryToSimplePolicyTest) {
        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        NRetry::TRetryPolicy policy{ MakeHolder<NSimple::TSimplePolicy>() };


        // Situation when algorithm finds backend instantly through slave policy
        // --------------------------------------------------------------------------------
        EXPECT_CALL(algoMock, SelectNext())
            .Times(1)
            .WillOnce(Return(&backMock));

        auto* backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------


        // Situation when algorithm finds any backend only after second call
        // So algorithm will be reseted after first failure
        // --------------------------------------------------------------------------------
        EXPECT_CALL(algoMock, SelectNext())
            .Times(2)
            .WillOnce(Return(nullptr))
            .WillOnce(Return(&backMock));

        EXPECT_CALL(algoMock, Reset())
            .Times(1);

        backend = policy.Next(&algoMock, true);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(UniqueRetryPolicyTest) {
        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        NUniqueRetry::TUniqueRetryOptions opts{};
        NUniqueRetry::TUniqueRetryPolicy policy(opts);


        // Situation when algorithm returns valid backend instantly
        // --------------------------------------------------------------------------------
        EXPECT_CALL(algoMock, Next())
            .Times(1)
            .WillOnce(Return(&backMock));

        EXPECT_CALL(algoMock, Reset())
            .Times(0);

        EXPECT_CALL(algoMock, RemoveSelected(&backMock))
            .Times(1);

        auto* backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------


        // Situation when algorithm can't find any backend
        // --------------------------------------------------------------------------------
        EXPECT_CALL(algoMock, Next())
            .Times(2)
            .WillRepeatedly(Return(nullptr));

        EXPECT_CALL(algoMock, Reset())
            .Times(1);

        EXPECT_CALL(algoMock, RemoveSelected(&backMock))
            .Times(0);

        backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, nullptr);
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(TimeoutPolicyTest) {
        constexpr auto TIMEOUT_SAMPLE = TDuration::Seconds(1);

        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        auto policyPtr = MakeHolder<NSrvKernel::NTesting::TPolicyMock>();
        auto* policyMock = policyPtr.Get();

        NTimeout::TTimeoutPolicy policy{ std::move(policyPtr), TIMEOUT_SAMPLE };


        // First call to Next sets the deadline
        // Second call to Next must be before that deadline
        // All other calls to Next must be after deadline so they will return nullptr
        // --------------------------------------------------------------------------------
        ON_CALL(*policyMock, Next(&algoMock, _))
            .WillByDefault(Return(&backMock));

        auto* backend = policy.Next(&algoMock, true);
        UNIT_ASSERT_EQUAL(backend, &backMock);

        backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, &backMock);

        Sleep(TIMEOUT_SAMPLE + TDuration::Seconds(0.5));

        for (auto i : xrange(ITERATIONS_NUM)) {
            Y_UNUSED(i);
            backend = policy.Next(&algoMock, true);
            UNIT_ASSERT_EQUAL(backend, nullptr);
        }
        // --------------------------------------------------------------------------------

        // Basic invariants
        // --------------------------------------------------------------------------------
        EXPECT_CALL(*policyMock, RegisterSuccess())
            .Times(1);

        policy.RegisterSuccess();

        EXPECT_CALL(*policyMock, RegisterFail())
            .Times(1);

        policy.RegisterFail();
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(TimeoutPolicyWithoutTimeoutTest) {
        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        auto policyPtr = MakeHolder<NSrvKernel::NTesting::TPolicyMock>();
        auto* policyMock = policyPtr.Get();

        NTimeout::TTimeoutPolicy policy{ std::move(policyPtr), TDuration{} };


        // When we don't supply timeout to policy, it just calls Next() in algorithm
        // --------------------------------------------------------------------------------
        ON_CALL(*policyMock, Next(&algoMock, _))
            .WillByDefault(Return(&backMock));

        for (auto i : xrange(ITERATIONS_NUM)) {
            Y_UNUSED(i);
            auto* backend = policy.Next(&algoMock, true);
            UNIT_ASSERT_EQUAL(backend, &backMock);
        }
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(ActivePolicyTest) {
        constexpr size_t ATTEMPTS_NUM = 10;

        NSrvKernel::NTesting::TAlgorithmMock algoMock;
        NSrvKernel::NTesting::TPerWorkerBackendMock backMock(MakeHolder<TBackendDescriptor>("", nullptr));

        auto policyPtr = MakeHolder<NSrvKernel::NTesting::TPolicyMock>();
        auto* policyMock = policyPtr.Get();

        NActive::TActivePolicy policy{ std::move(policyPtr), ATTEMPTS_NUM };


        // Situation when first choosen backend is enabled
        // No attempts used
        // --------------------------------------------------------------------------------
        backMock.SetEnabled(true);

        EXPECT_CALL(*policyMock, Next(&algoMock, _))
            .Times(1)
            .WillOnce(Return(&backMock));

        auto* backend = policy.Next(&algoMock, true);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------


        // Situation when first choosen backend is disabled and second backend wasn't found
        // One attempt used
        // --------------------------------------------------------------------------------
        backMock.SetEnabled(false);

        EXPECT_CALL(*policyMock, Next(&algoMock, _))
            .Times(2)
            .WillOnce(Return(&backMock))
            .WillOnce(Return(nullptr));

        backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, nullptr);
        // --------------------------------------------------------------------------------


        // Situation when all backends are disabled
        // All attempts used
        // --------------------------------------------------------------------------------
        backMock.SetEnabled(false);

        EXPECT_CALL(*policyMock, Next(&algoMock, _))
            .Times(ATTEMPTS_NUM)
            .WillRepeatedly(Return(&backMock));

        backend = policy.Next(&algoMock, true);
        UNIT_ASSERT_EQUAL(backend, nullptr);
        // --------------------------------------------------------------------------------


        // Even when all attempts are exhausted we firstly trying to get backend from slave
        // --------------------------------------------------------------------------------
        backMock.SetEnabled(true);

        EXPECT_CALL(*policyMock, Next(&algoMock, _))
            .Times(1)
            .WillOnce(Return(&backMock));

        backend = policy.Next(&algoMock, false);
        UNIT_ASSERT_EQUAL(backend, &backMock);
        // --------------------------------------------------------------------------------
    }
}
