package com.tisawesomeness.minecord.util.concurrent;

import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;

import javax.annotation.Nullable;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;

/**
 * A utility class for adding callbacks when a {@link CompletableFuture} completes.
 */
@Slf4j
public final class FutureCallback {
    private FutureCallback() {}

    /**
     * Creates a new callback builder. No changes will be applied until {@link Builder#build()} is called.
     * @param future The future to have callbacks added to
     * @param <T> The type of the future
     * @return A new builder
     */
    public static <T> Builder<T> builder(@NonNull CompletableFuture<T> future) {
        return new Builder<>(future);
    }

    /**
     * Builds success and error callbacks.
     * @param <T> The type of the provided future
     * @implNote This class does not violate the future's thread-safety,
     * but this class itself is <b>not</b> thread-safe.
     */
    public static class Builder<T> {
        private final @NonNull CompletableFuture<T> future;
        private @Nullable Consumer<? super T> successFn;
        private @Nullable Consumer<? super Throwable> errorFn;
        private @Nullable Consumer<? super Throwable> uncaughtFn;

        private Builder(@NonNull CompletableFuture<T> future) {
            this.future = future;
        }

        /**
         * Sets the success callback.
         * @param successFn The success function to run when the future succeeds
         * @return This builder
         */
        public Builder<T> onSuccess(@Nullable Consumer<? super T> successFn) {
            this.successFn = successFn;
            return this;
        }
        /**
         * Sets the failure callback.
         * @param errorFn The error function to run when the future completes exceptionally
         * @return This builder
         */
        public Builder<T> onFailure(@Nullable Consumer<? super Throwable> errorFn) {
            this.errorFn = errorFn;
            return this;
        }
        /**
         * Sets the uncaught callback. Defaults to simply logging the exception.
         * @param uncaughtFn The function to run when the success or error function throw an uncaught exception
         * @return This builder
         */
        public Builder<T> onUncaught(@Nullable Consumer<? super Throwable> uncaughtFn) {
            this.uncaughtFn = uncaughtFn;
            return this;
        }

        /**
         * Adds this builder's callbacks to the provided future.
         * @return The CompletableFuture with added callbacks
         */
        public CompletableFuture<T> build() {
            return future.whenComplete(this::onComplete);
        }
        private void onComplete(T val, Throwable err) {
            if (err == null) {
                if (successFn != null) {
                    trySuccess(val);
                }
            } else if (errorFn != null) {
                tryError(err);
            }
        }
        private void trySuccess(T val) {
            try {
                successFn.accept(val);
            } catch (Throwable ex) {
                if (uncaughtFn != null) {
                    uncaughtFn.accept(ex);
                } else {
                    log.error("Uncaught exception in success callback", ex);
                }
            }
        }
        private void tryError(Throwable err) {
            try {
                errorFn.accept(err);
            } catch (Throwable ex) {
                if (uncaughtFn != null) {
                    uncaughtFn.accept(ex);
                } else {
                    log.error("Uncaught exception in error callback", ex);
                }
            }
        }
    }

}

@@ -1289,16 +1289,63 @@
 aughtFn;
+%0A        private @Nullable Runnable completeFn;
 %0A%0A      
@@ -2524,32 +2524,388 @@
  this;%0A        %7D
+%0A        /**%0A         * Sets the completion callback.%0A         * @param completeFn The function to run when the success or error function throw an uncaught exception%0A         * @return This builder%0A         */%0A        public Builder%3CT%3E onComplete(@Nullable Runnable completeFn) %7B%0A            this.completeFn = completeFn;%0A            return this;%0A        %7D
 %0A%0A        /**%0A  
@@ -3203,32 +3203,115 @@
 hrowable err) %7B%0A
+            if (completeFn != null) %7B%0A                tryComplete();%0A            %7D%0A
             if (
@@ -3510,32 +3510,388 @@
     %7D%0A        %7D%0A
+        private void tryComplete() %7B%0A            try %7B%0A                completeFn.run();%0A            %7D catch (Throwable ex) %7B%0A                if (uncaughtFn != null) %7B%0A                    uncaughtFn.accept(ex);%0A                %7D else %7B%0A                    log.error(%22Uncaught exception in complete callback%22, ex);%0A                %7D%0A            %7D%0A        %7D%0A
         private 
@@ -4586,32 +4586,32 @@
 %7D%0A            %7D%0A
-
         %7D%0A    %7D%0A
@@ -4600,21 +4600,22 @@
 %7D%0A        %7D%0A
+%0A
     %7D%0A%0A%7D%0A

Add completeFn to callbacks