package cn.axzo.msg.center.utils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * @author yanglin
 */
public class AsyncTasks<T> {

    private final List<CompletableFuture<T>> futures = Collections.synchronizedList(new ArrayList<>());
    private final boolean removeFutureWhenComplete;

    public AsyncTasks() {
        this(true);
    }

    AsyncTasks(boolean removeFutureWhenComplete) {
        this.removeFutureWhenComplete = removeFutureWhenComplete;
    }

    public void add(CompletableFuture<T> future) {
        futures.add(future);
        if (removeFutureWhenComplete)
            future.whenComplete((unused, throwable) -> futures.remove(future));
    }

    CompletableFuture<T> getFuture(int index) {
        return futures.get(index);
    }

    public void awaitTermination() {
        collectAll().join();
    }

    public void awaitTermination(long timeout, TimeUnit unit)
            throws InterruptedException, ExecutionException, TimeoutException {
        collectAll().get(timeout, unit);
    }

    private CompletableFuture<Void> collectAll() {
        return CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]));
    }
}