package cn.reghao.tnb.gateway.limiter;

import java.util.Map;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.context.annotation.Primary;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setResponseStatus;

/**
 * 继承 org.springframework.cloud.gateway.filter.factory.RequestRateLimiterGatewayFilterFactory
 *
 * @author reghao
 * @date 2025-07-16 10:02:14
 */
@Slf4j
@Primary
@Component
public class MyRequestRateLimiterGatewayFilterFactory extends RequestRateLimiterGatewayFilterFactory {
    private static final String EMPTY_KEY = "____EMPTY_KEY__";
    private final ObjectMapper objectMapper;

    public MyRequestRateLimiterGatewayFilterFactory(RateLimiter defaultRateLimiter,
                                                    KeyResolver defaultKeyResolver,
                                                    ObjectMapper objectMapper) {
        super(defaultRateLimiter, defaultKeyResolver);
        this.objectMapper = objectMapper;
    }

    @Override
    public GatewayFilter apply(RequestRateLimiterGatewayFilterFactory.Config config) {
        KeyResolver resolver = getOrDefault(config.getKeyResolver(), super.getDefaultKeyResolver());
        RateLimiter<Object> limiter = getOrDefault(config.getRateLimiter(), super.getDefaultRateLimiter());
        boolean denyEmpty = getOrDefault(config.getDenyEmptyKey(), super.isDenyEmptyKey());
        HttpStatusHolder emptyKeyStatus = HttpStatusHolder.parse(getOrDefault(config.getEmptyKeyStatus(), super.getEmptyKeyStatusCode()));
        
        return (exchange, chain) -> resolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY).flatMap(key -> {
            if (EMPTY_KEY.equals(key)) {
                if (denyEmpty) {
                    setResponseStatus(exchange, emptyKeyStatus);
                    return exchange.getResponse().setComplete();
                }
                return chain.filter(exchange);
            }
            String routeId = config.getRouteId();
            if (routeId == null) {
                Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);
                assert route != null;
                routeId = route.getId();
            }

            return limiter.isAllowed(routeId, key).flatMap(response -> {
                for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                    exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
                }
                if (response.isAllowed()) {
                    return chain.filter(exchange);
                }

                ServerHttpResponse serverHttpResponse = exchange.getResponse();
                boolean ret = serverHttpResponse.setStatusCode(config.getStatusCode());
                if (!ret && log.isWarnEnabled()) {
                    log.warn("Unable to set status code to " + serverHttpResponse + ". Response already committed.");
                }
                serverHttpResponse.getHeaders().setContentType(MediaType.APPLICATION_JSON);
                return serverHttpResponse.writeWith(Mono.create(monoSink -> {
                    try {
                        String body = "REQUEST_RATE_LIMIT";
                        byte[] bytes = objectMapper.writeValueAsBytes(body);
                        DataBuffer dataBuffer = serverHttpResponse.bufferFactory().wrap(bytes);
                        monoSink.success(dataBuffer);
                    } catch (JsonProcessingException e) {
                        log.error(e.getMessage());
                        monoSink.error(e);
                    }
                }));
            });
        });
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return (configValue != null) ? configValue : defaultValue;
    }
}
