#include "rate_limiter.h"

#include "gtest/gtest.h"

#include "tile/base/chrono.h"
#include "tile/base/make_unique.h"
#include "tile/base/random.h"

namespace tile {

TEST(RateLimiter, TokenBucketRateLimiter)
{
    TokenBucketRateLimiter limiter(1000, 1);
    std::size_t total = 0;
    auto start        = ReadSteadyClock();
    while (ReadSteadyClock() - start < std::chrono::seconds(5)) {
        auto current = limiter.GetQuota();
        total += current;
        limiter.ConsumeBytes(current);

        // with sleep
        std::this_thread::sleep_for(std::chrono::milliseconds(1) * Random(100));
    }
    ASSERT_NEAR(6000, total, 200);
}

TEST(RateLimiter, TokenBucketRateLimiter2)
{
    TokenBucketRateLimiter limiter(1000, 1);
    std::size_t total = 0;
    auto start        = ReadSteadyClock();
    while (ReadSteadyClock() - start < std::chrono::seconds(5)) {
        auto current = limiter.GetQuota();
        total += current;
        limiter.ConsumeBytes(current);
    }
    ASSERT_NEAR(6000, total, 200);
}

TEST(RateLimiter, TokenBucketRateLimiterCapBurst)
{
    TokenBucketRateLimiter limiter(25, 500);
    for (int i = 0; i != 10; ++i) {
        ASSERT_EQ(25, limiter.GetQuota());
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }
}

TEST(RateLimiter, TokenBucketRateLimiterCapBurst2)
{
    TokenBucketRateLimiter limiter(1000, 500);
    for (int i = 0; i != 10; ++i) {
        ASSERT_EQ(1000, limiter.GetQuota());
        std::this_thread::sleep_for(std::chrono::milliseconds(10));// Enough to fully fill the bucket.
    }
}

TEST(RateLimiter, MultithreadedRateLimiter)
{
    ThreadSafeRateLimiter limiter(make_unique<TokenBucketRateLimiter>(1000, 1));
    std::atomic<std::size_t> total{0};
    std::vector<std::thread> ts;

    for (int i = 0; i != 10; ++i) {
        ts.emplace_back(std::thread([&] {
            auto start = ReadSteadyClock();
            while (ReadSteadyClock() - start < std::chrono::seconds(5)) {
                auto current = limiter.GetQuota();
                total += current;
                limiter.ConsumeBytes(current);
                std::this_thread::sleep_for(std::chrono::milliseconds(10) * Random(10));
            }
        }));
    }

    for (auto &&t : ts) { t.join(); }

    ASSERT_NEAR(6000, total.load(), 500);
}

TEST(RateLimiter, LayeredRateLimiter)
{
    ThreadSafeRateLimiter base_limiter(make_unique<TokenBucketRateLimiter>(1000, 1));
    auto our_limiter = make_unique<ThreadSafeRateLimiter>(make_unique<TokenBucketRateLimiter>(1000, 100));
    LayeredRateLimiter layered_limiter(&base_limiter, std::move(our_limiter));
    std::atomic<std::size_t> total{0};
    std::vector<std::thread> ts;

    for (int i = 0; i != 10; ++i) {
        ts.emplace_back(std::thread([&] {
            auto start = ReadSteadyClock();
            while (ReadSteadyClock() - start < std::chrono::seconds(5)) {
                auto current = layered_limiter.GetQuota();
                total += current;
                layered_limiter.ConsumeBytes(current);
                std::this_thread::sleep_for(std::chrono::milliseconds(10) * Random(10));
            }
        }));
    }

    for (auto &&t : ts) { t.join(); }

    ASSERT_NEAR(6000, total.load(), 500);// `msrl` takes effect.
}

TEST(RateLimiter, LayeredRateLimiter2)
{
    ThreadSafeRateLimiter base_limiter(make_unique<TokenBucketRateLimiter>(1000, 100));
    auto our_limiter = make_unique<ThreadSafeRateLimiter>(make_unique<TokenBucketRateLimiter>(1000, 1));
    LayeredRateLimiter layered_limiter(&base_limiter, std::move(our_limiter));
    std::atomic<std::size_t> total{0};
    std::vector<std::thread> ts;

    for (int i = 0; i != 10; ++i) {
        ts.emplace_back(std::thread([&] {
            auto start = ReadSteadyClock();
            while (ReadSteadyClock() - start < std::chrono::seconds(5)) {
                auto current = layered_limiter.GetQuota();
                total += current;
                layered_limiter.ConsumeBytes(current);
                std::this_thread::sleep_for(std::chrono::milliseconds(1) * Random(10));
            }
        }));
    }

    for (auto &&t : ts) { t.join(); }

    ASSERT_NEAR(6000, total.load(), 500);// `tbsrl` takes effect.
}

TEST(RateLimiter, LayeredRateLimiter3)
{
    ThreadSafeRateLimiter base_limiter(make_unique<TokenBucketRateLimiter>(1000, 1));
    std::atomic<std::size_t> total{0};
    std::vector<std::thread> ts;

    for (int i = 0; i != 10; ++i) {
        ts.emplace_back(std::thread([&] {
            auto our_limiter = make_unique<TokenBucketRateLimiter>(1000, 100);
            LayeredRateLimiter layered_limiter(&base_limiter, std::move(our_limiter));
            auto start = ReadSteadyClock();
            while (ReadSteadyClock() - start < std::chrono::seconds(5)) {
                auto current = layered_limiter.GetQuota();
                total += current;
                layered_limiter.ConsumeBytes(current);
            }
        }));
    }

    for (auto &&t : ts) { t.join(); }

    ASSERT_NEAR(6000, total.load(), 500);// `msrl` takes effect.
}

TEST(RateLimiter, LayeredRateLimiter4)
{
    ThreadSafeRateLimiter base_limiter(make_unique<TokenBucketRateLimiter>(1000, 100));
    std::atomic<std::size_t> total{0};
    std::vector<std::thread> ts;

    for (int i = 0; i != 10; ++i) {
        ts.emplace_back(std::thread([&] {
            auto our_limiter = make_unique<ThreadSafeRateLimiter>(make_unique<TokenBucketRateLimiter>(1000, 1));
            LayeredRateLimiter layered_limiter(&base_limiter, std::move(our_limiter));
            auto start = ReadSteadyClock();
            while (ReadSteadyClock() - start < std::chrono::seconds(5)) {
                auto current = layered_limiter.GetQuota();
                total += current;
                layered_limiter.ConsumeBytes(current);
            }
        }));
    }

    for (auto &&t : ts) { t.join(); }

    ASSERT_NEAR(60000, total.load(), 5000);// `tbsrl` takes effect.
}

}// namespace tile
