/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.apache.bookkeeper.common.concurrent;

import com.google.common.base.Stopwatch;
import com.google.common.collect.Lists;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.common.stats.OpStatsListener;
import org.apache.bookkeeper.common.util.OrderedScheduler;
import org.apache.bookkeeper.stats.OpStatsLogger;

/**
 * Future related utils.
 */
@Slf4j
public final class FutureUtils {

    private FutureUtils() {}

    private static final Function<Throwable, Exception> DEFAULT_EXCEPTION_HANDLER = cause -> {
        if (cause instanceof Exception) {
            return (Exception) cause;
        } else {
            return new Exception(cause);
        }
    };

    public static CompletableFuture<Void> Void() {
        return value(null);
    }

    public static <T> T result(CompletableFuture<T> future) throws Exception {
        return FutureUtils.result(future, DEFAULT_EXCEPTION_HANDLER);
    }

    public static <T> T result(CompletableFuture<T> future, long timeout, TimeUnit timeUnit) throws Exception {
        return FutureUtils.result(future, DEFAULT_EXCEPTION_HANDLER, timeout, timeUnit);
    }

    @SneakyThrows(InterruptedException.class)
    public static <T, ExceptionT extends Throwable> T result(
        CompletableFuture<T> future, Function<Throwable, ExceptionT> exceptionHandler) throws ExceptionT {
        try {
            try {
                /*
                 * CompletableFuture.get() in JDK8 spins before blocking and wastes CPU time.
                 * CompletableFuture.get(long, TimeUnit) blocks immediately (if the result is
                 * not yet available). While the implementation of get() has changed in JDK9
                 * (not spinning any more), using CompletableFuture.get(long, TimeUnit) allows
                 * us to avoid spinning for all current JDK versions.
                 */
                return future.get(Long.MAX_VALUE, TimeUnit.NANOSECONDS);
            } catch (TimeoutException eignore) {
                // it's ok to return null if we timeout after 292 years (2^63 nanos)
                return null;
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw e;
        } catch (ExecutionException e) {
            ExceptionT cause = exceptionHandler.apply(e.getCause());
            if (null == cause) {
                return null;
            } else {
                throw cause;
            }
        }
    }

    @SneakyThrows(InterruptedException.class)
    public static <T, ExceptionT extends Throwable> T result(
        CompletableFuture<T> future,
        Function<Throwable, ExceptionT> exceptionHandler,
        long timeout,
        TimeUnit timeUnit) throws ExceptionT, TimeoutException {
        try {
            return future.get(timeout, timeUnit);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw e;
        } catch (ExecutionException e) {
            ExceptionT cause = exceptionHandler.apply(e.getCause());
            if (null == cause) {
                return null;
            } else {
                throw cause;
            }
        }
    }

    public static <T> CompletableFuture<T> createFuture() {
        return new CompletableFuture<T>();
    }

    public static <T> CompletableFuture<T> value(T value) {
        return CompletableFuture.completedFuture(value);
    }

    public static <T> CompletableFuture<T> exception(Throwable cause) {
        CompletableFuture<T> future = FutureUtils.createFuture();
        future.completeExceptionally(cause);
        return future;
    }

    public static <T> void complete(CompletableFuture<T> result,
                                    T value) {
        if (null == result) {
            return;
        }
        result.complete(value);
    }

    public static <T> void completeExceptionally(CompletableFuture<T> result,
                                                 Throwable cause) {
        if (null == result) {
            return;
        }
        result.completeExceptionally(cause);
    }

    /**
     * Completing the {@code future} in the thread in the scheduler identified by
     * the {@code scheduleKey}.
     *
     * @param future      future to complete
     * @param action      action to execute when complete
     * @param scheduler   scheduler to execute the action.
     * @param scheduleKey key to choose the thread to execute the action
     * @param <T>
     * @return
     */
    public static <T> CompletableFuture<T> whenCompleteAsync(
        CompletableFuture<T> future,
        BiConsumer<? super T, ? super Throwable> action,
        OrderedScheduler scheduler,
        Object scheduleKey) {
        return future.whenCompleteAsync(action, scheduler.chooseThread(scheduleKey));
    }

    public static <T> CompletableFuture<List<T>> collect(List<CompletableFuture<T>> futureList) {
        CompletableFuture<Void> finalFuture =
            CompletableFuture.allOf(futureList.toArray(new CompletableFuture[futureList.size()]));
        return finalFuture.thenApply(result ->
            futureList
                .stream()
                .map(CompletableFuture::join)
                .collect(Collectors.toList()));
    }

    public static <T> void proxyTo(CompletableFuture<T> src,
                                   CompletableFuture<T> target) {
        src.whenComplete((value, cause) -> {
            if (null == cause) {
                target.complete(value);
            } else {
                target.completeExceptionally(cause);
            }
        });
    }

    //
    // Process futures
    //

    private static class ListFutureProcessor<T, R>
        implements FutureEventListener<R>, Runnable {

        private volatile boolean done = false;
        private final Iterator<T> itemsIter;
        private final Function<T, CompletableFuture<R>> processFunc;
        private final CompletableFuture<List<R>> promise;
        private final List<R> results;
        private final ExecutorService callbackExecutor;

        ListFutureProcessor(List<T> items,
                            Function<T, CompletableFuture<R>> processFunc,
                            ExecutorService callbackExecutor) {
            this.itemsIter = items.iterator();
            this.processFunc = processFunc;
            this.promise = new CompletableFuture<>();
            this.results = Lists.newArrayListWithExpectedSize(items.size());
            this.callbackExecutor = callbackExecutor;
        }

        @Override
        public void onSuccess(R value) {
            results.add(value);
            if (null == callbackExecutor) {
                run();
            } else {
                callbackExecutor.submit(this);
            }
        }

        @Override
        public void onFailure(final Throwable cause) {
            done = true;

            if (null == callbackExecutor) {
                promise.completeExceptionally(cause);
            } else {
                callbackExecutor.submit((Runnable) () -> promise.completeExceptionally(cause));
            }
        }

        @Override
        public void run() {
            if (done) {
                if (log.isDebugEnabled()) {
                    log.debug("ListFutureProcessor is interrupted.");
                }
                return;
            }
            if (!itemsIter.hasNext()) {
                promise.complete(results);
                done = true;
                return;
            }
            processFunc.apply(itemsIter.next()).whenComplete(this);
        }
    }

    /**
     * Process the list of items one by one using the process function <i>processFunc</i>.
     * The process will be stopped immediately if it fails on processing any one.
     *
     * @param collection       list of items
     * @param processFunc      process function
     * @param callbackExecutor executor to process the item
     * @return future presents the list of processed results
     */
    public static <T, R> CompletableFuture<List<R>> processList(List<T> collection,
                                                                Function<T, CompletableFuture<R>> processFunc,
                                                                @Nullable ExecutorService callbackExecutor) {
        ListFutureProcessor<T, R> processor =
            new ListFutureProcessor<T, R>(collection, processFunc, callbackExecutor);
        if (null != callbackExecutor) {
            callbackExecutor.submit(processor);
        } else {
            processor.run();
        }
        return processor.promise;
    }

    /**
     * Raise an exception to the <i>promise</i> within a given <i>timeout</i> period.
     * If the promise has been satisfied before raising, it won't change the state of the promise.
     *
     * @param promise   promise to raise exception
     * @param timeout   timeout period
     * @param unit      timeout period unit
     * @param cause     cause to raise
     * @param scheduler scheduler to execute raising exception
     * @param key       the submit key used by the scheduler
     * @return the promise applied with the raise logic
     */
    public static <T> CompletableFuture<T> within(final CompletableFuture<T> promise,
                                                  final long timeout,
                                                  final TimeUnit unit,
                                                  final Throwable cause,
                                                  final OrderedScheduler scheduler,
                                                  final Object key) {
        if (timeout < 0 || promise.isDone()) {
            return promise;
        }
        // schedule a timeout to raise timeout exception
        final java.util.concurrent.ScheduledFuture<?> task = scheduler.scheduleOrdered(key, () -> {
            if (!promise.isDone() && promise.completeExceptionally(cause)) {
                log.info("Raise exception", cause);
            }
        }, timeout, unit);
        // when the promise is satisfied, cancel the timeout task
        promise.whenComplete((value, throwable) -> {
                if (!task.cancel(true) && log.isDebugEnabled()) {
                    log.debug("Failed to cancel the timeout task");
                }
            }
        );
        return promise;
    }

    /**
     * Ignore exception from the <i>future</i>.
     *
     * @param future the original future
     * @return a transformed future ignores exceptions
     */
    public static <T> CompletableFuture<Void> ignore(CompletableFuture<T> future) {
        return ignore(future, null);
    }

    /**
     * Ignore exception from the <i>future</i> and log <i>errorMsg</i> on exceptions.
     *
     * @param future   the original future
     * @param errorMsg the error message to log on exceptions
     * @return a transformed future ignores exceptions
     */
    public static <T> CompletableFuture<Void> ignore(CompletableFuture<T> future,
                                                     final String errorMsg) {
        final CompletableFuture<Void> promise = new CompletableFuture<Void>();
        future.whenComplete(new FutureEventListener<T>() {
            @Override
            public void onSuccess(T value) {
                promise.complete(null);
            }

            @Override
            public void onFailure(Throwable cause) {
                if (null != errorMsg) {
                    log.error(errorMsg, cause);
                }
                promise.complete(null);
            }
        });
        return promise;
    }

    public static <T> CompletableFuture<T> ensure(CompletableFuture<T> future,
                                                  Runnable ensureBlock) {
        return future.whenComplete((value, cause) -> {
            ensureBlock.run();
        });
    }

    public static <T> CompletableFuture<T> rescue(CompletableFuture<T> future,
                                                  Function<Throwable, CompletableFuture<T>> rescueFuc) {
        CompletableFuture<T> result = FutureUtils.createFuture();
        future.whenComplete((value, cause) -> {
            if (null == cause) {
                result.complete(value);
                return;
            }
            proxyTo(rescueFuc.apply(cause), result);
        });
        return result;
    }

    /**
      * Add a event listener over <i>result</i> for collecting the operation stats.
      *
      * @param result result to listen on
      * @param opStatsLogger stats logger to record operations stats
      * @param stopwatch stop watch to time operation
      * @param <T>
      * @return result after registered the event listener
      */
    public static <T> CompletableFuture<T> stats(CompletableFuture<T> result,
                                                 OpStatsLogger opStatsLogger,
                                                 Stopwatch stopwatch) {
        return result.whenComplete(new OpStatsListener<T>(opStatsLogger, stopwatch));
    }

}
