package com.facebook.presto.hive;

import com.google.common.collect.ImmutableList;
import io.airlift.log.Logger;
import io.airlift.units.Duration;

import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

public class RetryDriver
{
    private static final Logger log = Logger.get(RetryDriver.class);
    private static final int DEFAULT_RETRY_ATTEMPTS = 10;
    private static final Duration DEFAULT_SLEEP_TIME = Duration.valueOf("1s");

    private final int maxRetryAttempts;
    private final Duration sleepTime;
    private final List<Class<? extends Exception>> exceptionWhiteList;

    private RetryDriver(int maxRetryAttempts, Duration sleepTime, List<Class<? extends Exception>> exceptionWhiteList)
    {
        this.maxRetryAttempts = maxRetryAttempts;
        this.sleepTime = sleepTime;
        this.exceptionWhiteList = exceptionWhiteList;

    }

    private RetryDriver()
    {
        this(DEFAULT_RETRY_ATTEMPTS, DEFAULT_SLEEP_TIME, ImmutableList.<Class<? extends Exception>>of());
    }

    public static RetryDriver retry()
    {
        return new RetryDriver();
    }

    public RetryDriver withMaxRetries(int maxRetryAttempts)
    {
        checkArgument(maxRetryAttempts > 0, "maxRetryAttempts must be greater than zero");
        return new RetryDriver(maxRetryAttempts, sleepTime, exceptionWhiteList);
    }

    public RetryDriver withSleep(Duration sleepTime)
    {
        return new RetryDriver(maxRetryAttempts, checkNotNull(sleepTime, "sleepTime is null"), exceptionWhiteList);
    }

    @SafeVarargs
    public final RetryDriver stopOn(Class<? extends Exception>... classes)
    {
        checkNotNull(classes, "classes is null");
        List<Class<? extends Exception>> exceptions = ImmutableList.<Class<? extends Exception>>builder()
                .addAll(exceptionWhiteList)
                .addAll(Arrays.asList(classes))
                .build();


        return new RetryDriver(maxRetryAttempts, sleepTime, exceptions);
    }

    public RetryDriver stopOnIllegalExceptions()
    {
        return stopOn(NullPointerException.class, IllegalStateException.class, IllegalArgumentException.class);
    }

    public <V> V run(String callableName, Callable<V> callable)
            throws Exception
    {
        checkNotNull(callableName, "callableName is null");
        checkNotNull(callable, "callable is null");

        int attempt = 0;
        while (true) {
            attempt++;
            try {
                return callable.call();
            }
            catch (Exception e) {
                for (Class<? extends Exception> clazz : exceptionWhiteList) {
                    if (clazz.isInstance(e)) {
                        throw e;
                    }
                }
                if (attempt >= maxRetryAttempts) {
                    throw e;
                }
                else {
                    log.debug("Failed on executing %s with attempt %d, will retry. Exception: %s", callableName, attempt, e.getMessage());
                }
                TimeUnit.MILLISECONDS.sleep((long) sleepTime.toMillis());
            }
        }
    }
}

@@ -630,16 +630,100 @@
 f(%221s%22);
+%0A    private static final Duration DEFAULT_MAX_RETRY_TIME = Duration.valueOf(%2230s%22);
 %0A%0A    pr
@@ -790,24 +790,65 @@
  sleepTime;%0A
+    private final Duration maxRetryTime;%0A
     private 
@@ -972,16 +972,39 @@
 eepTime,
+ Duration maxRetryTime,
  List%3CCl
@@ -1137,24 +1137,66 @@
  sleepTime;%0A
+        this.maxRetryTime = maxRetryTime;%0A
         this
@@ -1333,16 +1333,40 @@
 EP_TIME,
+ DEFAULT_MAX_RETRY_TIME,
  Immutab
@@ -1707,32 +1707,46 @@
 mpts, sleepTime,
+ maxRetryTime,
  exceptionWhiteL
@@ -1911,16 +1911,239 @@
  null%22),
+ maxRetryTime, exceptionWhiteList);%0A    %7D%0A%0A    public RetryDriver withMaxRetryTime(Duration maxRetryTime)%0A    %7B%0A        return new RetryDriver(maxRetryAttempts, sleepTime, checkNotNull(maxRetryTime, %22maxRetryTime is null%22),
  excepti
@@ -2592,24 +2592,38 @@
 , sleepTime,
+ maxRetryTime,
  exceptions)
@@ -3013,24 +3013,68 @@
 is null%22);%0A%0A
+        long startTime = System.nanoTime();%0A
         int 
@@ -3483,16 +3483,79 @@
 Attempts
+ %7C%7C Duration.nanosSince(startTime).compareTo(maxRetryTime) %3E= 0
 ) %7B%0A    

Add max retry time to RetryDriver