package ai.verta.modeldb.common.futures;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

public class InternalFuture<T> {
  private CompletionStage<T> stage;

  private InternalFuture() {}

  // Convert a list of futures to a future of a list
  public static <T> InternalFuture<List<T>> sequence(final List<InternalFuture<T>> futures, Executor executor) {
    final CompletableFuture<List<T>> promise = new CompletableFuture<>();
    final ArrayList<T> values = new ArrayList<>(futures.size());
    executor.execute(
            () -> {
              try {
                for (final var future : futures) {
                  values.add(future.stage.toCompletableFuture().join());
                }
                promise.complete(values);
              } catch (Throwable t) {
                promise.completeExceptionally(t);
              }
            });
    return InternalFuture.from(promise);
  }

  public static <R> InternalFuture<R> from(CompletionStage<R> other) {
    var ret = new InternalFuture<R>();
    ret.stage = other;
    return ret;
  }

  public static <R> InternalFuture<R> completedInternalFuture(R value) {
    var ret = new InternalFuture<R>();
    ret.stage = CompletableFuture.completedFuture(value);
    return ret;
  }

  public <U> InternalFuture<U> thenCompose(
      Function<? super T, InternalFuture<U>> fn, Executor executor) {
    return from(
        stage.thenComposeAsync(fn.andThen(internalFuture -> internalFuture.stage), executor));
  }

  public <U> InternalFuture<U> thenApply(Function<? super T, ? extends U> fn, Executor executor) {
    return from(stage.thenApplyAsync(fn, executor));
  }

  public InternalFuture<Void> thenAccept(Consumer<? super T> action, Executor executor) {
    return from(stage.thenAcceptAsync(action, executor));
  }

  public <U> InternalFuture<Void> thenRun(Runnable runnable, Executor executor) {
    return from(stage.thenRunAsync(runnable, executor));
  }

  public static <U> InternalFuture<Void> runAsync(Runnable runnable, Executor executor) {
    return from(CompletableFuture.runAsync(runnable, executor));
  }

  public static <U> InternalFuture<U> supplyAsync(Supplier<U> supplier, Executor executor) {
    return from(CompletableFuture.supplyAsync(supplier, executor));
  }

  public static <U> InternalFuture<U> failedStage(Throwable ex) {
    return from(CompletableFuture.failedFuture(ex));
  }

  public InternalFuture<T> whenComplete(
      BiConsumer<? super T, ? super Throwable> action, Executor executor) {
    return from(stage.whenCompleteAsync(action, executor));
  }

  public T get() throws ExecutionException, InterruptedException {
    return stage.toCompletableFuture().get();
  }

  public CompletionStage<T> toCompletionStage() {
    return stage;
  }
}

@@ -654,16 +654,23 @@
 equence(
+%0A      
 final Li
@@ -731,16 +731,34 @@
   final 
+var promise = new 
 Completa
@@ -779,94 +779,294 @@
 %3CT%3E%3E
- promise = new CompletableFuture%3C%3E
+();%0A    final var values = new ArrayList%3CT%3E(futures.size());%0A    final CompletableFuture%3CT%3E%5B%5D futuresArray =%0A        futures.stream
 ()
-;
 %0A    
-final ArrayList%3CT%3E values = new ArrayList%3C%3E(
+        .map(x -%3E x.toCompletionStage().toCompletableFuture())%0A            .collect(Collectors.toList())%0A            .toArray(new CompletableFuture%5B
 futu
@@ -1079,31 +1079,76 @@
 ze()
+%5D
 );%0A
+%0A
     
-executor.execute
+CompletableFuture.allOf(futuresArray)%0A        .thenRunAsync
 (%0A  
@@ -1232,16 +1232,21 @@
  futures
+Array
 ) %7B%0A    
@@ -1281,40 +1281,11 @@
 ure.
-stage.toCompletableFuture().join
+get
 ());
@@ -1462,19 +1462,42 @@
        %7D
+,%0A            executor
 );%0A
+%0A
     retu

Fix sequence to avoid deadlocks (#2122)