#include "tile/io/util/rate_limiter.h"

#include "tile/base/chrono.h"
#include "tile/base/logging.h"

namespace tile {

namespace {
class NullLimiter : public RateLimiter {
public:
    std::size_t GetQuota() override { return 0; }

    void ConsumeBytes(std::size_t consumed) override {}
};
}// namespace

RateLimiter *
RateLimiter::GetDefaultRxRateLimiter()
{
    static NullLimiter null_limiter;
    return &null_limiter;
}

RateLimiter *
RateLimiter::GetDefaultTxRateLimiter()
{
    static NullLimiter null_limiter;
    return &null_limiter;
}

TokenBucketRateLimiter::TokenBucketRateLimiter(std::size_t bucket_quota,
                                               std::size_t quota_per_tick,
                                               std::chrono::nanoseconds tick,
                                               bool over_consumption_allowed)
    : max_quota_(bucket_quota),
      quota_per_tick_(quota_per_tick),
      tick_(tick),
      over_consumption_allowed_(over_consumption_allowed)
{
    TILE_CHECK_GT(bucket_quota, 0);
    TILE_CHECK_GT(quota_per_tick, 0);
    last_refill_ = ReadSteadyClock().time_since_epoch() / tick_;
    curr_quota_  = max_quota_;
}

std::size_t
TokenBucketRateLimiter::GetQuota()
{
    auto now                  = ReadSteadyClock().time_since_epoch() / tick_;
    std::uint64_t last_refill = internal::Exchange(last_refill_, now);

    curr_quota_ += quota_per_tick_ * (now - last_refill);

    if (curr_quota_ > 0) {
        curr_quota_ = std::min<std::size_t>(curr_quota_, max_quota_);
        return curr_quota_;
    } else {
        return 0;
    }
}

void
TokenBucketRateLimiter::ConsumeBytes(std::size_t consumed)
{
    TILE_CHECK(over_consumption_allowed_ || consumed <= curr_quota_);
    curr_quota_ -= consumed;
}

ThreadSafeRateLimiter::ThreadSafeRateLimiter(MaybeOwning<RateLimiter> limiter, std::size_t burst_limit)
    : burst_limit_(burst_limit),
      impl_(std::move(limiter))
{
    TILE_CHECK_GT(burst_limit_, 0);
}

std::size_t
ThreadSafeRateLimiter::GetQuota()
{
    std::lock_guard<std::mutex> lock(lock_);
    return impl_->GetQuota();
}

void
ThreadSafeRateLimiter::ConsumeBytes(std::size_t consumed)
{
    std::lock_guard<std::mutex> lock(lock_);
    impl_->ConsumeBytes(consumed);
}

LayeredRateLimiter::LayeredRateLimiter(RateLimiter *upper, MaybeOwning<RateLimiter> ours)
    : upper_(upper),
      ours_(std::move(ours))
{}

std::size_t
LayeredRateLimiter::GetQuota()
{
    return std::min(upper_->GetQuota(), ours_->GetQuota());
}

void
LayeredRateLimiter::ConsumeBytes(std::size_t consumed)
{
    upper_->ConsumeBytes(consumed);
    ours_->ConsumeBytes(consumed);
}
}// namespace tile
