/**
 * 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.filter;

import java.util.Map;

import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.support.HasRouteId;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

import com.crunii.micro.gateway.limit.GuavaBurstyRateLimiter;
import com.crunii.micro.gateway.limit.resolver.PathKeyResolver;

/**
 * @author 田平 create 2020年1月19日下午5:37:49
 *
 */
@Component
public class GuavaRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<GuavaRateLimiterGatewayFilterFactory.Config> {

	private static final String EMPTY_KEY = "_GuavaRateLimiter_EMPTY_KEY_";

	private final RateLimiter<?> defaultRateLimiter = new GuavaBurstyRateLimiter();

	private final KeyResolver defaultKeyResolver = new PathKeyResolver();

	private boolean denyEmptyKey = true;

	private String emptyKeyStatusCode = HttpStatus.FORBIDDEN.name();

	public GuavaRateLimiterGatewayFilterFactory() {
		super(Config.class);
	}

	public KeyResolver getDefaultKeyResolver() {
		return defaultKeyResolver;
	}

	public RateLimiter<?> getDefaultRateLimiter() {
		return defaultRateLimiter;
	}

	public boolean isDenyEmptyKey() {
		return denyEmptyKey;
	}

	public void setDenyEmptyKey(boolean denyEmptyKey) {
		this.denyEmptyKey = denyEmptyKey;
	}

	public String getEmptyKeyStatusCode() {
		return emptyKeyStatusCode;
	}

	public void setEmptyKeyStatusCode(String emptyKeyStatusCode) {
		this.emptyKeyStatusCode = emptyKeyStatusCode;
	}

	@Override
	public GatewayFilter apply(Config config) {
		KeyResolver resolver = getOrDefault(config.keyResolver, defaultKeyResolver);
		RateLimiter<?> limiter = getOrDefault(config.rateLimiter, defaultRateLimiter);
		boolean denyEmpty = getOrDefault(config.denyEmptyKey, this.denyEmptyKey);
		HttpStatusHolder emptyKeyStatus = HttpStatusHolder.parse(getOrDefault(config.emptyKeyStatus, this.emptyKeyStatusCode));

		return (exchange, chain) -> resolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY)
				.flatMap(key -> {
					if (EMPTY_KEY.equals(key)) {
						if (denyEmpty) {
							ServerWebExchangeUtils.setResponseStatus(exchange, emptyKeyStatus);
							return exchange.getResponse().setComplete();
						}
						return chain.filter(exchange);
					}
					return limiter.isAllowed(config.getRouteId(), 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);
								}
								ServerWebExchangeUtils.setResponseStatus(exchange, config.getStatusCode());
								return exchange.getResponse().setComplete();
							});
				});
	}

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

	public static class Config implements HasRouteId {

		private KeyResolver keyResolver;

		private RateLimiter<?> rateLimiter;

		private HttpStatus statusCode = HttpStatus.TOO_MANY_REQUESTS;

		private Boolean denyEmptyKey;

		private String emptyKeyStatus;

		private String routeId;

		public KeyResolver getKeyResolver() {
			return keyResolver;
		}

		public Config setKeyResolver(KeyResolver keyResolver) {
			this.keyResolver = keyResolver;
			return this;
		}

		public RateLimiter<?> getRateLimiter() {
			return rateLimiter;
		}

		public Config setRateLimiter(RateLimiter<?> rateLimiter) {
			this.rateLimiter = rateLimiter;
			return this;
		}

		public HttpStatus getStatusCode() {
			return statusCode;
		}

		public Config setStatusCode(HttpStatus statusCode) {
			this.statusCode = statusCode;
			return this;
		}

		public Boolean getDenyEmptyKey() {
			return denyEmptyKey;
		}

		public Config setDenyEmptyKey(Boolean denyEmptyKey) {
			this.denyEmptyKey = denyEmptyKey;
			return this;
		}

		public String getEmptyKeyStatus() {
			return emptyKeyStatus;
		}

		public Config setEmptyKeyStatus(String emptyKeyStatus) {
			this.emptyKeyStatus = emptyKeyStatus;
			return this;
		}

		@Override
		public void setRouteId(String routeId) {
			this.routeId = routeId;
		}

		@Override
		public String getRouteId() {
			return this.routeId;
		}

	}
}
