package cn.reghao.tnb.sb.ratelimiter;

import lombok.Getter;
import lombok.Setter;

import java.time.LocalTime;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 滑动时间窗口算法
 *
 * @author reghao
 * @date 2025-07-15 13:52:29
 */
public class SlidingWindowRateLimiter {
    // 阈值
    private final int qps;
    // 时间窗口(毫秒)
    private long timeWindow = 1000;
    // 子窗口数量(每 100ms 滑动一次窗口)
    private Integer windowCount = 10;
    // 窗口列表
    private WindowInfo[] windowArray = new WindowInfo[windowCount];

    public SlidingWindowRateLimiter(int qps) {
        this.qps = qps;
        long currentTimeMillis = System.currentTimeMillis();
        for (int i = 0; i < windowArray.length; i++) {
            windowArray[i] = new WindowInfo(currentTimeMillis, new AtomicInteger(0));
        }
    }

    public synchronized boolean tryAcquire() {
        long currentTimeMillis = System.currentTimeMillis();
        // 1.计算当前时间窗口位置
        int currentIndex = (int)(currentTimeMillis % timeWindow / (timeWindow / windowCount));
        // 2.更新当前窗口计数 & 重置过期窗口计数
        int sum = 0;
        for (int i = 0; i < windowArray.length; i++) {
            WindowInfo windowInfo = windowArray[i];
            if ((currentTimeMillis - windowInfo.getTime()) > timeWindow) {
                // 重置过期窗口计数
                windowInfo.getNumber().set(0);
                windowInfo.setTime(currentTimeMillis);
            }
            if (currentIndex == i && windowInfo.getNumber().get() < qps) {
                // 更新当前窗口计数
                windowInfo.getNumber().incrementAndGet();
            }
            sum = sum + windowInfo.getNumber().get();
        }
        // 3.判断当前 QPS 是否超过限制
        return sum <= qps;
    }

    @Setter
    @Getter
    static class WindowInfo {
        // 窗口开始时间
        private Long time;
        // 计数器
        private AtomicInteger number;

        public WindowInfo(long time, AtomicInteger number) {
            this.time = time;
            this.number = number;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        int qps = 2, count = 20, sleep = 300, success = count*sleep/1000*qps;
        System.out.printf("当前 QPS 阈值: %d, 当前请求次数: %d, 请求间隔: %dms, 预计成功次数:%d\n", qps, count, sleep, success);
        success = 0;
        SlidingWindowRateLimiter myRateLimiter = new SlidingWindowRateLimiter(qps);
        for (int i = 0; i < count; i++) {
            Thread.sleep(sleep);
            if (myRateLimiter.tryAcquire()) {
                success++;
                System.out.println(LocalTime.now() + ": success");
            } else {
                System.out.println(LocalTime.now() + ": fail");
            }
        }

        System.out.printf("实际测试成功次数: %s\n", success);
    }
}
