Nextus-Web/src/main/java/dev/coph/nextusweb/server/ratelimit/SlidingWindowLimiter.java

package dev.coph.nextusweb.server.ratelimit;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

public final class SlidingWindowLimiter implements RateLimiter {

    private final long limit;
    private final long windowNanos;
    private final ConcurrentHashMap<String, SlidingWindow> windows = new ConcurrentHashMap<>();

    public SlidingWindowLimiter(long limit, long windowMillis) {
        this.limit = limit;
        this.windowNanos = windowMillis * 1_000_000L;
    }

    @Override
    public Result tryAcquire(String key, long nowNanos) {
        SlidingWindow w = windows.computeIfAbsent(key, k -> new SlidingWindow(nowNanos));
        return w.tryAcquire(nowNanos, limit, windowNanos);
    }

    public void cleanup(long olderThanNanos) {
        long now = System.nanoTime();
        windows.entrySet().removeIf(e -> now - e.getValue().windowStart.get() > olderThanNanos);
    }

    private static final class SlidingWindow {
        final AtomicLong windowStart;
        final AtomicLong currentCount;
        final AtomicLong previousCount;

        SlidingWindow(long now) {
            this.windowStart = new AtomicLong(now);
            this.currentCount = new AtomicLong(0);
            this.previousCount = new AtomicLong(0);
        }

        synchronized Result tryAcquire(long now, long limit, long windowNanos) {
            long start = windowStart.get();
            long elapsed = now - start;

            if (elapsed >= 2 * windowNanos) {
                windowStart.set(now);
                previousCount.set(0);
                currentCount.set(0);
                elapsed = 0;
            } else if (elapsed >= windowNanos) {
                windowStart.set(start + windowNanos);
                previousCount.set(currentCount.get());
                currentCount.set(0);
                elapsed -= windowNanos;
            }

            double prevWeight = 1.0 - ((double) elapsed / windowNanos);
            long weightedCount = (long) (previousCount.get() * prevWeight) + currentCount.get();

            if (weightedCount >= limit) {
                long retryMs = (windowNanos - elapsed) / 1_000_000L;
                return Result.deny(limit, Math.max(1, retryMs));
            }

            currentCount.incrementAndGet();
            return Result.allow(limit - weightedCount - 1, limit);
        }
    }
}