package org.infinispan.interceptors.impl;

import java.util.concurrent.CompletableFuture;
import java.util.function.BiConsumer;

import org.infinispan.commands.VisitableCommand;
import org.infinispan.context.InvocationContext;
import org.infinispan.interceptors.InvocationCallback;
import org.infinispan.util.concurrent.CompletableFutures;
import org.infinispan.util.logging.Log;
import org.infinispan.util.logging.LogFactory;
import org.jgroups.annotations.GuardedBy;

/**
 * Invocation stage representing a computation that may or may not be done yet.
 *
 * <p>It stores handler objects in a queue instead of creating a new instance every time a handler is added.
 * The queue may be frozen based on internal conditions, like executing the last handler or reaching the capacity
 * of the queue, and adding a handler will create a new instance.
 * The queue will also be frozen when {@link #toCompletableFuture()} is invoked, to make that future behave like
 * a regular {@link CompletableFuture}.
 * </p>
 *
 * <p>When the queue is not frozen, adding a handler will change the result of the current stage.
 * When the queue is frozen, adding a handler may actually execute the handler synchronously.</p>
 *
 * @author Dan Berindei
 * @since 9.0
 */
public class QueueAsyncInvocationStage extends SimpleAsyncInvocationStage implements BiConsumer<Object, Throwable>,
      InvocationCallback {
   private static final Log log = LogFactory.getLog(QueueAsyncInvocationStage.class);
   private static final boolean trace = log.isTraceEnabled();

   private final InvocationContext ctx;
   private final VisitableCommand command;

   public QueueAsyncInvocationStage(InvocationContext ctx, VisitableCommand command,
                                    CompletableFuture<?> valueFuture, InvocationCallback function) {
      super(new CompletableFuture<>());
      this.ctx = ctx;
      this.command = command;

      queueAdd(function);
      valueFuture.whenComplete(this);
   }

   @Override
   public Object addCallback(InvocationContext ctx, VisitableCommand command, InvocationCallback function) {
      if (ctx != this.ctx || command != this.command) {
         return new SimpleAsyncInvocationStage(future).addCallback(ctx, command, function);
      }

      if (queueAdd(function)) {
         return this;
      }

      return invokeDirectly(ctx, command, function);
   }

   private Object invokeDirectly(InvocationContext ctx, VisitableCommand command, InvocationCallback function) {
      Object rv;
      Throwable throwable;
      try {
         rv = future.join();
         throwable = null;
      } catch (Throwable t) {
         rv = null;
         throwable = CompletableFutures.extractException(t);
      }
      try {
         return function.apply(ctx, command, rv, throwable);
      } catch (Throwable t) {
         return new SimpleAsyncInvocationStage(t);
      }
   }

   @Override
   public void accept(Object rv, Throwable throwable) {
      // We started with a CompletableFuture, which is now complete.
      invokeQueuedHandlers(rv, throwable);
   }

   @Override
   public Object apply(InvocationContext rCtx, VisitableCommand rCommand, Object rv, Throwable throwable)
         throws Throwable {
      // We started from another AsyncInvocationStage, which is now complete.
      invokeQueuedHandlers(rv, throwable);

      if (throwable == null) {
         return rv;
      } else {
         throw throwable;
      }
   }

   private void invokeQueuedHandlers(Object rv, Throwable throwable) {
      if (trace)
         log.tracef("Resuming invocation of command %s with %d handlers", command, queueSize());
      while (true) {
         InvocationCallback function = queuePoll();
         if (function == null) {
            // Complete the future.
            // We finished running the handlers, and the last pollHandler() call locked the queue.
            if (throwable == null) {
               future.complete(rv);
            } else {
               future.completeExceptionally(throwable);
            }
            return;
         }

         // Run the handler
         try {
            if (throwable != null) {
               throwable = CompletableFutures.extractException(throwable);
            }
            rv = function.apply(ctx, command, rv, throwable);
            throwable = null;
         } catch (Throwable t) {
            rv = null;
            throwable = t;
         }
         if (rv instanceof SimpleAsyncInvocationStage) {
            SimpleAsyncInvocationStage currentStage = (SimpleAsyncInvocationStage) rv;
            if (!currentStage.isDone()) {
               if (currentStage instanceof QueueAsyncInvocationStage) {
                  QueueAsyncInvocationStage queueAsyncInvocationStage = (QueueAsyncInvocationStage) currentStage;
                  queueAsyncInvocationStage.future.whenComplete(this);
                  return;
               } else {
                  // Use the CompletableFuture directly, without creating another AsyncInvocationStage instance
                  currentStage.future.whenComplete(this);
                  return;
               }
            } else {
               try {
                  rv = currentStage.get();
               } catch (Throwable t) {
                  throwable = t;
               }
            }
         }
         // We got a synchronous invocation stage, continue with the next handler
      }
   }


   /**
    * An inline implementation of a queue. The queue is frozen automatically when the last element is polled.
    */
   // Capacity must be a power of 2
   static final int QUEUE_INITIAL_CAPACITY = 8;
   @GuardedBy("this")
   private InvocationCallback[] elements = null;
   @GuardedBy("this")
   private byte mask;
   @GuardedBy("this")
   private byte head;
   @GuardedBy("this")
   private byte tail;

   @GuardedBy("this")
   private boolean frozen;

   boolean queueAdd(InvocationCallback element) {
      synchronized (this) {
         if (frozen) {
            return false;
         }
         if (elements == null || tail - head > mask) {
            queueExpand();
         }
         elements[tail & mask] = element;
         tail++;
         return true;
      }
   }

   /**
    * Remove one handler from the deque, or freeze the deque if there are no more elements.
    *
    * @return The next handler, or {@code null} if the deque is empty.
    */
   InvocationCallback queuePoll() {
      InvocationCallback element;
      synchronized (this) {
         if (tail != head) {
            element = elements[head & mask];
            head++;
         } else {
            element = null;
            frozen = true;
         }
      }
      return element;
   }

   /**
    * @return The current number of elements in the deque, only useful for debugging.
    */
   int queueSize() {
      synchronized (this) {
         // We can assume tail and head won't overflow in our use case
         return tail - head;
      }
   }

   @GuardedBy("this")
   private void queueExpand() {
      // We start with no elements and mask 0
      if (elements == null) {
         elements = new InvocationCallback[QUEUE_INITIAL_CAPACITY];
         mask = QUEUE_INITIAL_CAPACITY - 1;
         return;
      }

      InvocationCallback[] oldElements = elements;
      int oldCapacity = oldElements.length;
      int oldMask = mask;
      int oldHead = head;
      int oldTail = tail;
      int maskedHead = oldHead & oldMask;
      int maskedTail = oldTail & oldMask;
      int oldSize = tail - head;
      if (oldSize != oldCapacity)
         throw new IllegalStateException("Queue should be expanded only when full");

      int newSize = oldCapacity * 2;
      int newMask = newSize - 1;
      elements = new InvocationCallback[newSize];
      mask = (byte) newMask;
      head = 0;
      tail = (byte) (oldTail - oldHead);

      if (maskedHead < maskedTail) {
         System.arraycopy(oldElements, maskedHead, elements, 0, oldSize);
      } else {
         System.arraycopy(oldElements, maskedHead, elements, 0, oldCapacity - maskedHead);
         System.arraycopy(oldElements, 0, elements, oldCapacity - maskedHead, maskedTail);
      }
   }
}
