package com.cxyxiaomo.epp.gateway.Factory;

import com.google.common.util.concurrent.RateLimiter;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
// 一个自定义的限流过滤器工厂
public class RateLimitByIpGatewayFilterFactory extends AbstractGatewayFilterFactory<RateLimitByIpGatewayFilterFactory.Config> {

    // 用于存储IP地址和对应的计数器
    private static final Map<String, RateLimiter> RATE_LIMITER_CACHE = new ConcurrentHashMap<>();

    public RateLimitByIpGatewayFilterFactory() {
        super(Config.class);
    }

    @Override
    public GatewayFilter apply(Config config) {
        return (exchange, chain) -> {
            // 获取请求的IP地址
            ServerHttpRequest request = exchange.getRequest();
            String ip = request.getRemoteAddress().getAddress().getHostAddress();

            // 根据IP地址获取对应的限流器
            RateLimiter rateLimiter = RATE_LIMITER_CACHE.get(ip);
            if (rateLimiter == null) {
                // 如果不存在，则创建一个新的限流器，并放入缓存中
                rateLimiter = RateLimiter.create(config.getRate());
                RATE_LIMITER_CACHE.put(ip, rateLimiter);
            }

            // 判断请求是否被限流
            if (rateLimiter.tryAcquire(config.getPermits())) {
                // 如果没有被限流，则放行
                return chain.filter(exchange);
            } else {
                System.out.println("限流！ip: " + ip);
                // 如果被限流，则返回429状态码（Too Many Requests）
                ServerHttpResponse response = exchange.getResponse();
                response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                return response.setComplete();
            }
        };
    }

    // 配置类，用于接收配置参数
    public static class Config {

        // 每秒允许的请求数
        private double rate;

        // 每次请求需要的令牌数
        private int permits;

        public double getRate() {
            return rate;
        }

        public void setRate(double rate) {
            this.rate = rate;
        }

        public int getPermits() {
            return permits;
        }

        public void setPermits(int permits) {
            this.permits = permits;
        }
    }
}
