/**
 * Copyright 2008-2009. Chongqing Communications Industry Services Co.,Ltd Information Technology Branch. All rights reserved. <a>http://www.crunii.com</a>
 */
package com.crunii.micro.gateway.limit;

import com.google.common.util.concurrent.RateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationService;
import reactor.core.publisher.Mono;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @author 田平 create 2020年1月19日下午3:46:02
 */
public class GuavaWarmingUpRateLimiter extends AbstractRateLimiter<GuavaWarmingUpRateLimiter.Config> {

	public static Map<String, RateLimiter> cacheLimiter = new ConcurrentHashMap<>();

	public static final String configPropName = "guava-warming-up-rate-limit";

	public GuavaWarmingUpRateLimiter() {
		super(Config.class, configPropName, null);
	}

	public GuavaWarmingUpRateLimiter(ConfigurationService configurationService) {
		super(Config.class, configPropName, configurationService);
	}

	@Override
	public Mono<Response> isAllowed(String routeId, String id) {
		RateLimiter limiter = cacheLimiter.computeIfAbsent(routeId, k -> {
			Config config = getConfig().get(routeId);
			return RateLimiter.create(config.getPermitsPerSecond(), config.getWarmupPeriod(), TimeUnit.SECONDS);
		});

		boolean isAllow = limiter.tryAcquire();
		HashMap<String, String> headers = new HashMap<String, String>();
		headers.put("limit_rate", String.valueOf(limiter.getRate()));
		return Mono.just(new Response(isAllow, headers));
	}

	public static class Config {
		// 限流默认每秒钟20个
		private Double permitsPerSecond = 20D;
		// 预热时间,默认1秒
		private int warmupPeriod = 1;// 单位是秒

		public int getWarmupPeriod() {
			return warmupPeriod;
		}

		public Double getPermitsPerSecond() {
			return permitsPerSecond;
		}

		public Config setPermitsPerSecond(Double permitsPerSecond) {
			this.permitsPerSecond = permitsPerSecond;
			return this;
		}

		public Config setPermitsPerSecond(int warmupPeriod) {
			this.warmupPeriod = warmupPeriod;
			return this;
		}
	}
}
