// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#include "token_bucket_rate_limiter.h"

#include <bthread/bthread.h>

#include <glog/logging.h> // IWYU pragma: export

#include <chrono>

#include <mutex>

#include <thread>

#if defined(__APPLE__)

#include <ctime>

#endif

namespace doris {

// Just 10^6.

static constexpr auto NS = 1000000000UL;

class TokenBucketRateLimiter::SimpleSpinLock {

public:

    SimpleSpinLock() = default;

    ~SimpleSpinLock() = default;

    void lock() {

        int spin_count = 0;

        static constexpr int MAX_SPIN_COUNT = 50;

        while (_flag.test_and_set(std::memory_order_acq_rel)) {

            spin_count++;

            if (spin_count >= MAX_SPIN_COUNT) {

                LOG(WARNING) << "Warning: Excessive spinning detected while acquiring lock. Spin "

                                "count: "

                             << spin_count;

                spin_count = 0;

            }

            // Spin until we acquire the lock

        }

    }

    void unlock() { _flag.clear(std::memory_order_release); }

private:

    std::atomic_flag _flag = ATOMIC_FLAG_INIT;

};

TokenBucketRateLimiter::TokenBucketRateLimiter(size_t max_speed, size_t max_burst, size_t limit)

        : _max_speed(max_speed),

          _max_burst(max_burst),

          _limit(limit),

          _mutex(std::make_unique<TokenBucketRateLimiter::SimpleSpinLock>()),

          _remain_tokens(max_burst) {}

TokenBucketRateLimiter::~TokenBucketRateLimiter() = default;

TokenBucketRateLimiterHolder::~TokenBucketRateLimiterHolder() = default;

std::pair<size_t, double> TokenBucketRateLimiter::_update_remain_token(long now, size_t amount) {

    // Values obtained under lock to be checked after release

    size_t count_value;

    double tokens_value;

    {

        std::lock_guard<SimpleSpinLock> lock(*_mutex);

        now = (now < _prev_ns_count) ? _prev_ns_count : now;

        if (_max_speed) {

            double delta_seconds =

                    _prev_ns_count ? static_cast<double>(now - _prev_ns_count) / NS : 0;

            _remain_tokens = std::min<double>(_remain_tokens + _max_speed * delta_seconds - amount,

                                              _max_burst);

        }

        _count += amount;

        count_value = _count;

        tokens_value = _remain_tokens;

        _prev_ns_count = now;

    }

    return {count_value, tokens_value};

}

int64_t TokenBucketRateLimiter::add(size_t amount) {

    // Values obtained under lock to be checked after release

    auto duration = std::chrono::steady_clock::now().time_since_epoch();

    auto time_nano_count = std::chrono::duration_cast<std::chrono::nanoseconds>(duration).count();

    auto [count_value, tokens_value] = _update_remain_token(time_nano_count, amount);

    if (_limit && count_value > _limit) {

        // CK would throw exception

        return -1;

    }

    // Wait unless there is positive amount of remain_tokens - throttling

    int64_t sleep_time_ns = 0;

    if (_max_speed && tokens_value < 0) {

        sleep_time_ns = static_cast<int64_t>(-tokens_value / _max_speed * NS);

        bthread_usleep(sleep_time_ns / 1000);

    }

    return sleep_time_ns;

}

TokenBucketRateLimiterHolder::TokenBucketRateLimiterHolder(size_t max_speed, size_t max_burst,

                                                           size_t limit,

                                                           std::function<void(int64_t)> metric_func)

        : rate_limiter(std::make_unique<TokenBucketRateLimiter>(max_speed, max_burst, limit)),

          metric_func(std::move(metric_func)) {}

int64_t TokenBucketRateLimiterHolder::add(size_t amount) {

    int64_t sleep;

    {

        std::shared_lock read {rate_limiter_rw_lock};

        sleep = rate_limiter->add(amount);

    }

    metric_func(sleep);

    return sleep;

}

int TokenBucketRateLimiterHolder::reset(size_t max_speed, size_t max_burst, size_t limit) {

    {

        std::unique_lock write {rate_limiter_rw_lock};

        rate_limiter = std::make_unique<TokenBucketRateLimiter>(max_speed, max_burst, limit);

    }

    return 0;

}

std::string to_string(S3RateLimitType type) {

    switch (type) {

    case S3RateLimitType::GET:

        return "get";

    case S3RateLimitType::PUT:

        return "put";

    default:

        return std::to_string(static_cast<size_t>(type));

    }

}

S3RateLimitType string_to_s3_rate_limit_type(std::string_view value) {

    if (value == "get") {

        return S3RateLimitType::GET;

    } else if (value == "put") {

        return S3RateLimitType::PUT;

    }

    return S3RateLimitType::UNKNOWN;

}

} // namespace doris