package cn.reghao.tnb.gateway.limiter;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationService;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

/**
 * 参照 org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter
 *
 * 已经存在一个默认实现 RedisRateLimiter, 若要使用自定义实现则需要使用 @Primary 注解
 *
 * @author reghao
 * @date 2025-07-16 10:02:14
 */
@Slf4j
@Primary
@Component
public class MyRateLimiter extends AbstractRateLimiter<MyRateLimiter.Config> {
    private static final String CONFIGURATION_PROPERTY_NAME = "my-rate-limiter";
    private RateLimiter rateLimiter = RateLimiter.create(10);

    protected MyRateLimiter(ConfigurationService configurationService) {
        super(MyRateLimiter.Config.class, CONFIGURATION_PROPERTY_NAME, configurationService);
    }

    @Override
    public Mono<Response> isAllowed(String routeId, String id) {
        Config config = getConfig().get(routeId);
        return Mono.fromSupplier(() -> {
            boolean acquire = rateLimiter.tryAcquire(config.requestedTokens);
            if (acquire) {
                return new Response(true, Maps.newHashMap());
            } else {
                log.info("{} refuse {} access", routeId, id);
                return new Response(false, Maps.newHashMap());
            }
        });
    }

    @Getter
    @Setter
    @ToString
    public static class Config {
        /**
         * 每次请求多少个 token
         */
        private Integer requestedTokens;
    }
}
