/*
 * Copyright 2013-2017 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.cloud.netflix.ribbon;

import java.util.ArrayList;
import java.util.List;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryContext;
import org.springframework.cloud.client.loadbalancer.LoadBalancedRetryPolicy;
import org.springframework.cloud.client.loadbalancer.ServiceInstanceChooser;
import org.springframework.http.HttpMethod;
import org.springframework.util.StringUtils;

import com.netflix.client.config.CommonClientConfigKey;
import com.netflix.client.config.IClientConfig;
import com.netflix.client.config.IClientConfigKey;

/**
 * {@link LoadBalancedRetryPolicy} for Ribbon clients.
 * @author Ryan Baxter
 */
public class RibbonLoadBalancedRetryPolicy implements LoadBalancedRetryPolicy {

	public static final IClientConfigKey<String> RETRYABLE_STATUS_CODES = new CommonClientConfigKey<String>("retryableStatusCodes") {};
	private static final Log log = LogFactory.getLog(RibbonLoadBalancedRetryPolicy.class);
	private int sameServerCount = 0;
	private int nextServerCount = 0;
	private String serviceId;
	private RibbonLoadBalancerContext lbContext;
	private ServiceInstanceChooser loadBalanceChooser;
	List<Integer> retryableStatusCodes = new ArrayList<>();

	public RibbonLoadBalancedRetryPolicy(String serviceId, RibbonLoadBalancerContext context, ServiceInstanceChooser loadBalanceChooser) {
		this.serviceId = serviceId;
		this.lbContext = context;
		this.loadBalanceChooser = loadBalanceChooser;
	}

	public RibbonLoadBalancedRetryPolicy(String serviceId, RibbonLoadBalancerContext context, ServiceInstanceChooser loadBalanceChooser,
										 IClientConfig clientConfig) {
		this.serviceId = serviceId;
		this.lbContext = context;
		this.loadBalanceChooser = loadBalanceChooser;
		String retryableStatusCodesProp = clientConfig.getPropertyAsString(RETRYABLE_STATUS_CODES, "");
		String[] retryableStatusCodesArray = retryableStatusCodesProp.split(",");
		for(String code : retryableStatusCodesArray) {
			if(!StringUtils.isEmpty(code)) {
				try {
					retryableStatusCodes.add(Integer.valueOf(code.trim()));
				} catch (NumberFormatException e) {
					log.warn("We cant add the status code because the code [ " + code + " ] could not be converted to an integer. ", e);
				}
			}
		}
	}

	public boolean canRetry(LoadBalancedRetryContext context) {
		HttpMethod method = context.getRequest().getMethod();
		return HttpMethod.GET == method || lbContext.isOkToRetryOnAllOperations();
	}

	@Override
	public boolean canRetrySameServer(LoadBalancedRetryContext context) {
		return sameServerCount < lbContext.getRetryHandler().getMaxRetriesOnSameServer() && canRetry(context);
	}

	@Override
	public boolean canRetryNextServer(LoadBalancedRetryContext context) {
		//this will be called after a failure occurs and we increment the counter
		//so we check that the count is less than or equals to too make sure
		//we try the next server the right number of times
		return nextServerCount <= lbContext.getRetryHandler().getMaxRetriesOnNextServer() && canRetry(context);
	}

	@Override
	public void close(LoadBalancedRetryContext context) {

	}

	@Override
	public void registerThrowable(LoadBalancedRetryContext context, Throwable throwable) {
		//Check if we need to ask the load balancer for a new server.
		//Do this before we increment the counters because the first call to this method
		//is not a retry it is just an initial failure.
		if(!canRetrySameServer(context)  && canRetryNextServer(context)) {
			context.setServiceInstance(loadBalanceChooser.choose(serviceId));
		}
		//This method is called regardless of whether we are retrying or making the first request.
		//Since we do not count the initial request in the retry count we don't reset the counter
		//until we actually equal the same server count limit.  This will allow us to make the initial
		//request plus the right number of retries.
		if(sameServerCount >= lbContext.getRetryHandler().getMaxRetriesOnSameServer() && canRetry(context)) {
			//reset same server since we are moving to a new server
			sameServerCount = 0;
			nextServerCount++;
			if(!canRetryNextServer(context)) {
				context.setExhaustedOnly();
			}
		} else {
			sameServerCount++;
		}

	}

	@Override
	public boolean retryableStatusCode(int statusCode) {
		return retryableStatusCodes.contains(statusCode);
	}
}
