#ifndef TILE_IO_UTIL_RATE_LIMITER_H
#define TILE_IO_UTIL_RATE_LIMITER_H

#pragma once

#include "tile/base/maybe_owning.h"

#include <limits>
#include <mutex>

namespace tile {
class RateLimiter {
public:
    virtual ~RateLimiter()                          = default;

    virtual std::size_t GetQuota()                  = 0;
    virtual void ConsumeBytes(std::size_t consumed) = 0;
    static RateLimiter *GetDefaultRxRateLimiter();
    static RateLimiter *GetDefaultTxRateLimiter();
};

class TokenBucketRateLimiter : public RateLimiter {
public:
    TokenBucketRateLimiter(std::size_t bucket_quota,
                           std::size_t quota_per_tick,
                           std::chrono::nanoseconds tick = std::chrono::milliseconds(1),
                           bool over_consumption_allowed = true);

    std::size_t GetQuota() override;
    void ConsumeBytes(std::size_t consumed) override;

private:
    std::size_t max_quota_;
    std::size_t quota_per_tick_;
    std::chrono::nanoseconds tick_;
    bool over_consumption_allowed_;

    std::uint64_t last_refill_;
    std::int64_t curr_quota_{0};
};

class ThreadSafeRateLimiter : public RateLimiter {
public:
    explicit ThreadSafeRateLimiter(MaybeOwning<RateLimiter> limiter,
                                   std::size_t burst_limit = std::numeric_limits<std::size_t>::max());
    std::size_t GetQuota() override;
    void ConsumeBytes(std::size_t consumed) override;

private:
    std::size_t burst_limit_;
    std::mutex lock_;
    MaybeOwning<RateLimiter> impl_;
};

class LayeredRateLimiter : public RateLimiter {
public:
    LayeredRateLimiter(RateLimiter *upper, MaybeOwning<RateLimiter> ours);

    std::size_t GetQuota() override;
    void ConsumeBytes(std::size_t consumed) override;

private:
    RateLimiter *upper_;
    MaybeOwning<RateLimiter> ours_;
};

}// namespace tile

#endif// TILE_IO_UTIL_RATE_LIMITER_H
