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.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;

/**
 * @author reghao
 * @date 2023-09-14 11:52:52
 */
@Slf4j
//@Primary
//@Component
public class MyRateLimiter extends AbstractRateLimiter<MyRateLimiter.Config> {
    /**
     * 和配置文件中的配置属性相对应
     */
    private static final String CONFIGURATION_PROPERTY_NAME = "default-gateway-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;
    }
}
