package cn.quantgroup.customer.config.rest;

import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpRequest;
import org.apache.http.annotation.Contract;
import org.apache.http.annotation.ThreadingBehavior;
import org.apache.http.client.protocol.HttpClientContext;
import org.apache.http.impl.client.DefaultHttpRequestRetryHandler;
import org.apache.http.protocol.HttpContext;
import org.apache.http.util.Args;

import javax.net.ssl.SSLException;
import java.io.IOException;
import java.net.UnknownHostException;
import java.util.Arrays;
import java.util.Locale;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * HTTP 自动重试;
 * 如果是post请求,并且异常有 Read timed out则不进行重试
 *
 * @author Jie.Feng
 * @date 2017/12/13
 */
@Contract(
        threading = ThreadingBehavior.IMMUTABLE
)
public class HttpRequestRetryHandler extends DefaultHttpRequestRetryHandler {
    private static final String POST = "POST";
    private static final String TIME_OUT = "Read timed out";
    private final Map<String, Boolean> idempotentMethods;

    /**
     * 现在重试的类型
     *
     * @param retryCount
     * @param requestSentRetryEnabled
     */
    public HttpRequestRetryHandler(int retryCount, boolean requestSentRetryEnabled) {
        super(retryCount, requestSentRetryEnabled, Arrays.asList(UnknownHostException.class, SSLException.class));
        this.idempotentMethods = new ConcurrentHashMap<>();
        this.idempotentMethods.put("GET", Boolean.TRUE);
        this.idempotentMethods.put("HEAD", Boolean.TRUE);
        this.idempotentMethods.put("PUT", Boolean.TRUE);
        this.idempotentMethods.put("DELETE", Boolean.TRUE);
        this.idempotentMethods.put("OPTIONS", Boolean.TRUE);
        this.idempotentMethods.put("TRACE", Boolean.TRUE);
        this.idempotentMethods.put("POST", Boolean.TRUE);
    }

    /**
     * 这算是默认所有都重3次
     */
    public HttpRequestRetryHandler() {
        this(3, false);
    }


    @Override
    protected boolean handleAsIdempotent(HttpRequest request) {
        String method = request.getRequestLine().getMethod().toUpperCase(Locale.ROOT);
        Boolean b = this.idempotentMethods.get(method);
        return b != null && b;
    }

    /**
     * 限制重试的类别
     *
     * @param exception
     * @param executionCount
     * @param context
     * @return
     */
    @Override
    public boolean retryRequest(IOException exception, int executionCount, HttpContext context) {
        Args.notNull(exception, "Exception parameter");
        Args.notNull(context, "HTTP context");
        final HttpClientContext clientContext = HttpClientContext.adapt(context);
        final HttpRequest request = clientContext.getRequest();
        String method = request.getRequestLine().getMethod().toUpperCase(Locale.ROOT);
        if (POST.equals(method)) {
            if (!StringUtils.containsIgnoreCase(exception.getMessage(), TIME_OUT)) {
                return super.retryRequest(exception, executionCount, context);
            } else {
                return false;
            }
        }
        return super.retryRequest(exception, executionCount, context);
    }
}
