package org.ta.gateway.config;

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.core.Ordered;
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.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 基于内存的本地限流过滤器（无 Redis 依赖，仅单机生效）
 * 实现令牌桶算法，适配 Reactive 环境
 */
@Component
public class LocalRateLimitFilter implements GatewayFilter, Ordered {

    // 限流缓存：key=客户端IP，value=令牌桶
    private final Map<String, TokenBucket> rateLimitCache = new ConcurrentHashMap<>();

    // 令牌桶参数：每秒生成 20 个令牌，最大容量 40
    private static final int REPLENISH_RATE = 20;
    private static final int BURST_CAPACITY = 40;
    private static final long INTERVAL = 1000 / REPLENISH_RATE; // 令牌生成间隔（毫秒）

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 1. 获取客户端 IP
        ServerHttpRequest request = exchange.getRequest();
        String clientIp = getClientIp(request);

        // 2. 获取/初始化令牌桶
        TokenBucket tokenBucket = rateLimitCache.computeIfAbsent(clientIp, k -> new TokenBucket(BURST_CAPACITY));

        // 3. 尝试获取令牌
        if (tokenBucket.tryAcquire()) {
            // 有令牌，继续请求
            return chain.filter(exchange);
        } else {
            // 无令牌，返回 429 Too Many Requests
            ServerHttpResponse response = exchange.getResponse();
            response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
            response.getHeaders().add("Content-Type", "application/json");
            String errorMsg = "{\"code\":429,\"message\":\"请求过于频繁，请稍后重试\"}";
            return response.writeWith(Mono.just(response.bufferFactory().wrap(errorMsg.getBytes())));
        }
    }

    /**
     * 获取客户端真实 IP
     */
    private String getClientIp(ServerHttpRequest request) {
        String ip = request.getHeaders().getFirst("X-Forwarded-For");
        if (ip == null || ip.isEmpty()) {
            ip = request.getRemoteAddress().getAddress().getHostAddress();
        }
        return ip;
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 10; // 优先级高于其他过滤器
    }

    /**
     * 基于内存的令牌桶实现（线程安全）
     */
    private static class TokenBucket {
        private final int capacity; // 令牌桶最大容量
        private final AtomicLong tokens; // 当前令牌数
        private volatile long lastRefillTime; // 上次填充令牌时间

        public TokenBucket(int capacity) {
            this.capacity = capacity;
            this.tokens = new AtomicLong(capacity);
            this.lastRefillTime = System.currentTimeMillis();
        }

        /**
         * 尝试获取 1 个令牌
         */
        public boolean tryAcquire() {
            // 1. 填充令牌（根据时间差补充）
            refillTokens();

            // 2. 尝试获取令牌
            long currentTokens = tokens.get();
            if (currentTokens > 0) {
                return tokens.compareAndSet(currentTokens, currentTokens - 1);
            }
            return false;
        }

        /**
         * 填充令牌（按时间间隔补充）
         */
        private void refillTokens() {
            long now = System.currentTimeMillis();
            long timeElapsed = now - lastRefillTime;

            // 计算应补充的令牌数
            long tokensToAdd = timeElapsed / INTERVAL;
            if (tokensToAdd > 0) {
                long newTokens = Math.min(capacity, tokens.get() + tokensToAdd);
                tokens.set(newTokens);
                lastRefillTime = now;
            }
        }
    }
}