package org.jeecg.config.GuavaRateLimiter;

import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Slf4j
public class GuavaRateLimiterUtils {
    private final static GuavaRateLimiterUtils INSTANCE = new GuavaRateLimiterUtils();

    private ConcurrentHashMap<String, RateLimiter> RATE_LIMITER_MAP = new ConcurrentHashMap<>();

    private GuavaRateLimiterUtils() {}

    /**
     * 单例
     * @return 返回单例对象
     */
    public static GuavaRateLimiterUtils getInstance() {
        return INSTANCE;
    }

    /**
     * 设置限流值
     * @param path 请求路径
     * @param permitsPerSecond 一秒内的限流数
     */
    public void setPermits(String path, Double permitsPerSecond) {
        if (!RATE_LIMITER_MAP.containsKey(path)) {
            RATE_LIMITER_MAP.put(path, RateLimiter.create(permitsPerSecond));
        }
    }

    /**
     * 尝试获取令牌
     * @param requestPath 接口路径
     * @param timeout 超时时间
     * @param timeUnit 超时时间单位
     * @return 获取是否成功
     */
    public Boolean tryAcquire(String requestPath, Long timeout, TimeUnit timeUnit) {
        if (!RATE_LIMITER_MAP.containsKey(requestPath)) {
            log.debug("限流列表中未查询到该路径：{}，直接放行", requestPath);
            return true;
        }
        RateLimiter rateLimiter = RATE_LIMITER_MAP.get(requestPath);
        if (!rateLimiter.tryAcquire(timeout, timeUnit)) {
            log.info("当前接口限流：{}", requestPath);
            return false;
        }
        return true;
    }

    /**
     * 第一次接口结束以后，根据本次执行时间，自适应调整QPS，使其尽量接近TPS
     * @param path 接口路径
     * @param spendTime 本次执行时间
     */
    public void adaptivePermits(String path, Long spendTime) {
        if (RATE_LIMITER_MAP.containsKey(path)) {
            RATE_LIMITER_MAP.get(path).setRate(1000.00 / spendTime);
            log.info("{}限流自适应调整成功：permits: {}", path, RATE_LIMITER_MAP.get(path).getRate());
        }
    }

}

