package cn.novalon.gym.manage.sys.filter;

import io.github.resilience4j.ratelimiter.RateLimiter;
import io.github.resilience4j.ratelimiter.RequestNotPermitted;
import io.github.resilience4j.ratelimiter.RateLimiterRegistry;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilter;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

import java.util.concurrent.ConcurrentHashMap;

/**
 * API限流过滤器
 * 
 * @author 张翔
 * @date 2026-03-13
 */
@Component
public class RateLimitFilter implements WebFilter {

    private final RateLimiter rateLimiter;
    private final ConcurrentHashMap<String, RateLimiter> ipRateLimiterMap = new ConcurrentHashMap<>();

    public RateLimitFilter(RateLimiterRegistry rateLimiterRegistry) {
        this.rateLimiter = rateLimiterRegistry.rateLimiter("apiRateLimiter");
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
        String clientIp = getClientIp(exchange);
        RateLimiter ipRateLimiter = ipRateLimiterMap.computeIfAbsent(clientIp, k -> rateLimiter);

        return Mono.fromCallable(() -> ipRateLimiter.acquirePermission())
                .flatMap(permitted -> {
                    if (permitted) {
                        return chain.filter(exchange);
                    } else {
                        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                        exchange.getResponse().getHeaders().add("X-RateLimit-Limit", 
                            String.valueOf(ipRateLimiter.getRateLimiterConfig().getLimitForPeriod()));
                        exchange.getResponse().getHeaders().add("X-RateLimit-Remaining", "0");
                        exchange.getResponse().getHeaders().add("Retry-After", "1");
                        return exchange.getResponse().setComplete();
                    }
                })
                .onErrorResume(RequestNotPermitted.class, e -> {
                    exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                    exchange.getResponse().getHeaders().add("X-RateLimit-Limit", 
                        String.valueOf(ipRateLimiter.getRateLimiterConfig().getLimitForPeriod()));
                    exchange.getResponse().getHeaders().add("X-RateLimit-Remaining", "0");
                    exchange.getResponse().getHeaders().add("Retry-After", "1");
                    return exchange.getResponse().setComplete();
                });
    }

    private String getClientIp(ServerWebExchange exchange) {
        String ip = exchange.getRequest().getHeaders().getFirst("X-Forwarded-For");
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = exchange.getRequest().getHeaders().getFirst("X-Real-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = exchange.getRequest().getRemoteAddress() != null 
                ? exchange.getRequest().getRemoteAddress().getAddress().getHostAddress() 
                : "unknown";
        }
        return ip;
    }
}