package cn.novalon.manage.gateway.filter;

import cn.novalon.manage.gateway.config.RateLimitConfig;
import io.github.resilience4j.ratelimiter.RateLimiter;
import io.github.resilience4j.ratelimiter.RequestNotPermitted;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.http.HttpHeaders;
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 org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

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

/**
 * 网关限流过滤器
 * 
 * 文件定义：实现多维度限流策略的全局过滤器
 * 涉及业务：API访问频率控制，防止滥用和DDoS攻击
 * 算法：使用Resilience4j的RateLimiter实现令牌桶算法
 * 
 * 限流维度：
 * 1. 全局限流：保护系统整体稳定性
 * 2. IP限流：防止单个IP过度访问
 * 3. 用户限流：防止单个用户过度访问
 * 
 * @author 张翔
 * @date 2026-03-26
 */
@Component
public class RateLimitFilter implements GlobalFilter, Ordered {

    private static final Logger logger = LoggerFactory.getLogger(RateLimitFilter.class);
    private static final String USER_ID_HEADER = "X-User-Id";
    private static final String RATE_LIMIT_REMAINING_HEADER = "X-RateLimit-Remaining";
    private static final String RATE_LIMIT_LIMIT_HEADER = "X-RateLimit-Limit";
    private static final String RETRY_AFTER_HEADER = "Retry-After";

    private final RateLimiter globalRateLimiter;
    private final RateLimiter ipRateLimiter;
    private final RateLimiter userRateLimiter;
    private final RateLimitConfig rateLimitConfig;

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

    private final AtomicInteger totalRequests = new AtomicInteger(0);
    private final AtomicInteger blockedRequests = new AtomicInteger(0);

    public RateLimitFilter(
            RateLimiter globalRateLimiter,
            RateLimiter ipRateLimiter,
            RateLimiter userRateLimiter,
            RateLimitConfig rateLimitConfig) {
        this.globalRateLimiter = globalRateLimiter;
        this.ipRateLimiter = ipRateLimiter;
        this.userRateLimiter = userRateLimiter;
        this.rateLimitConfig = rateLimitConfig;

        logger.info("RateLimitFilter initialized with enabled: {}", rateLimitConfig.isRateLimitEnabled());
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        if (!rateLimitConfig.isRateLimitEnabled()) {
            return chain.filter(exchange);
        }

        totalRequests.incrementAndGet();

        ServerHttpRequest request = exchange.getRequest();
        String clientIp = getClientIp(request);
        String userId = getUserId(request);
        String requestPath = request.getPath().value();

        logger.debug("Processing request - IP: {}, UserId: {}, Path: {}", clientIp, userId, requestPath);

        return checkGlobalRateLimit(exchange, chain, clientIp, userId);
    }

    private Mono<Void> checkGlobalRateLimit(
            ServerWebExchange exchange,
            GatewayFilterChain chain,
            String clientIp,
            String userId) {
        return Mono.fromCallable(() -> globalRateLimiter.acquirePermission())
                .flatMap(permitted -> {
                    if (permitted) {
                        return checkIpRateLimit(exchange, chain, clientIp, userId);
                    } else {
                        return handleRateLimitExceeded(exchange, "Global", clientIp, userId);
                    }
                })
                .onErrorResume(RequestNotPermitted.class,
                        e -> handleRateLimitExceeded(exchange, "Global", clientIp, userId));
    }

    private Mono<Void> checkIpRateLimit(
            ServerWebExchange exchange,
            GatewayFilterChain chain,
            String clientIp,
            String userId) {
        RateLimiter ipLimiter = ipRateLimiterMap.computeIfAbsent(
                clientIp,
                k -> createIpRateLimiter(clientIp));

        return Mono.fromCallable(() -> ipLimiter.acquirePermission())
                .flatMap(permitted -> {
                    if (permitted) {
                        if (userId != null && !userId.isEmpty()) {
                            return checkUserRateLimit(exchange, chain, userId);
                        } else {
                            return chain.filter(exchange);
                        }
                    } else {
                        return handleRateLimitExceeded(exchange, "IP", clientIp, userId);
                    }
                })
                .onErrorResume(RequestNotPermitted.class,
                        e -> handleRateLimitExceeded(exchange, "IP", clientIp, userId));
    }

    private Mono<Void> checkUserRateLimit(
            ServerWebExchange exchange,
            GatewayFilterChain chain,
            String userId) {
        RateLimiter userLimiter = userRateLimiterMap.computeIfAbsent(
                userId,
                k -> createUserRateLimiter(userId));

        return Mono.fromCallable(() -> userLimiter.acquirePermission())
                .flatMap(permitted -> {
                    if (permitted) {
                        return chain.filter(exchange);
                    } else {
                        return handleRateLimitExceeded(exchange, "User", null, userId);
                    }
                })
                .onErrorResume(RequestNotPermitted.class, e -> handleRateLimitExceeded(exchange, "User", null, userId));
    }

    private Mono<Void> handleRateLimitExceeded(
            ServerWebExchange exchange,
            String limitType,
            String clientIp,
            String userId) {
        blockedRequests.incrementAndGet();

        logger.warn("Rate limit exceeded - Type: {}, IP: {}, UserId: {}", limitType, clientIp, userId);

        ServerHttpResponse response = exchange.getResponse();
        response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);

        HttpHeaders headers = response.getHeaders();
        headers.add(RATE_LIMIT_LIMIT_HEADER, "0");
        headers.add(RATE_LIMIT_REMAINING_HEADER, "0");
        headers.add(RETRY_AFTER_HEADER, "1");
        headers.add("X-RateLimit-Type", limitType);

        String errorMessage = String.format(
                "{\"error\":\"Rate limit exceeded\",\"type\":\"%s\",\"message\":\"Too many requests. Please try again later.\"}",
                limitType);

        return response.writeWith(Mono.just(response.bufferFactory().wrap(errorMessage.getBytes())));
    }

    private String getClientIp(ServerHttpRequest request) {
        String ip = request.getHeaders().getFirst("X-Forwarded-For");
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeaders().getFirst("X-Real-IP");
        }
        if (ip == null || ip.isEmpty() || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddress() != null
                    ? request.getRemoteAddress().getAddress().getHostAddress()
                    : "unknown";
        }
        if (ip != null && ip.contains(",")) {
            ip = ip.split(",")[0].trim();
        }
        return ip;
    }

    private String getUserId(ServerHttpRequest request) {
        return request.getHeaders().getFirst(USER_ID_HEADER);
    }

    private RateLimiter createIpRateLimiter(String ip) {
        logger.debug("Creating rate limiter for IP: {}", ip);
        return RateLimiter.of("ip-" + ip, ipRateLimiter.getRateLimiterConfig());
    }

    private RateLimiter createUserRateLimiter(String userId) {
        logger.debug("Creating rate limiter for user: {}", userId);
        return RateLimiter.of("user-" + userId, userRateLimiter.getRateLimiterConfig());
    }

    public int getTotalRequests() {
        return totalRequests.get();
    }

    public int getBlockedRequests() {
        return blockedRequests.get();
    }

    public void resetCounters() {
        totalRequests.set(0);
        blockedRequests.set(0);
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 100;
    }
}
