package arez;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Objects;
import java.util.stream.Stream;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import static org.realityforge.braincheck.Guards.*;

/**
 * Basic implementation of task queue that supports priority based queuing of tasks.
 */
final class MultiPriorityTaskQueue
  implements TaskQueue
{
  /**
   * A buffer per priority containing tasks that have been scheduled but are not executing.
   */
  @Nonnull
  private final CircularBuffer<Task>[] _buffers;

  /**
   * Construct queue with specified priority count where each priority is backed by a buffer with specified size.
   *
   * @param priorityCount   the number of priorities supported.
   * @param initialCapacity the initial size of buffer for each priority.
   */
  @SuppressWarnings( { "unchecked", "rawtypes" } )
  MultiPriorityTaskQueue( final int priorityCount, final int initialCapacity )
  {
    assert priorityCount > 0;
    assert initialCapacity > 0;
    _buffers = (CircularBuffer<Task>[]) new CircularBuffer[ priorityCount ];
    for ( int i = 0; i < priorityCount; i++ )
    {
      _buffers[ i ] = new CircularBuffer<>( initialCapacity );
    }
  }

  /**
   * Return the number of priorities handled by the queue.
   *
   * @return the number of priorities handled by the queue.
   */
  int getPriorityCount()
  {
    return _buffers.length;
  }

  @Override
  public int getQueueSize()
  {
    int count = 0;
    //noinspection ForLoopReplaceableByForEach
    for ( int i = 0; i < _buffers.length; i++ )
    {
      count += _buffers[ i ].size();
    }
    return count;
  }

  @Override
  public boolean hasTasks()
  {
    //noinspection ForLoopReplaceableByForEach
    for ( int i = 0; i < _buffers.length; i++ )
    {
      if ( _buffers[ i ].isNotEmpty() )
      {
        return true;
      }
    }
    return false;
  }

  /**
   * Add the specified task to the queue.
   * The task must not already be in the queue.
   *
   * @param task the task.
   */
  @Override
  public void queueTask( @Nonnull final Task task )
  {
    queueTask( task.getPriorityIndex(), task );
  }

  /**
   * Add the specified task to the queue.
   * The task must not already be in the queue.
   *
   * @param priority the task priority.
   * @param task     the task.
   */
  void queueTask( final int priority, @Nonnull final Task task )
  {
    if ( Arez.shouldCheckInvariants() )
    {
      invariant( () -> Arrays.stream( _buffers ).noneMatch( b -> b.contains( task ) ),
                 () -> "Arez-0099: Attempting to queue task named '" + task.getName() +
                       "' when task is already queued." );
      invariant( () -> priority >= 0 && priority < _buffers.length,
                 () -> "Arez-0215: Attempting to queue task named '" + task.getName() +
                       "' but passed an invalid priority " + priority + "." );
    }
    Objects.requireNonNull( task ).markAsQueued();
    _buffers[ priority ].add( Objects.requireNonNull( task ) );
  }

  @Nullable
  @Override
  public Task dequeueTask()
  {
    // Return the highest priority taskQueue that has tasks in it and return task.
    //noinspection ForLoopReplaceableByForEach
    for ( int i = 0; i < _buffers.length; i++ )
    {
      final CircularBuffer<Task> taskQueue = _buffers[ i ];
      if ( taskQueue.isNotEmpty() )
      {
        final Task task = taskQueue.pop();
        assert null != task;
        return task;
      }
    }
    return null;
  }

  @Override
  public Collection<Task> clear()
  {
    final ArrayList<Task> tasks = new ArrayList<>();
    //noinspection ForLoopReplaceableByForEach
    for ( int i = 0; i < _buffers.length; i++ )
    {
      final CircularBuffer<Task> buffer = _buffers[ i ];
      Task task;
      while ( null != ( task = buffer.pop() ) )
      {
        tasks.add( task );
        task.markAsIdle();
      }
    }
    return tasks;
  }

  @Nonnull
  @Override
  public Stream<Task> getOrderedTasks()
  {
    assert Arez.shouldCheckInvariants() || Arez.shouldCheckApiInvariants();
    return Stream.of( _buffers ).flatMap( CircularBuffer::stream );
  }

  @Nonnull
  CircularBuffer<Task> getBufferByPriority( final int priority )
  {
    return _buffers[ priority ];
  }
}

@@ -1448,24 +1448,142 @@
 ength;%0A  %7D%0A%0A
+  /**%0A   * Return the number of tasks inside the queue.%0A   *%0A   * @return the number of tasks inside the queue.%0A   */%0A
   @Override%0A
@@ -1574,32 +1574,32 @@
  */%0A  @Override%0A
-
   public int get
@@ -1796,24 +1796,136 @@
 count;%0A  %7D%0A%0A
+  /**%0A   * Return true if queue has any tasks in it.%0A   *%0A   * @return true if queue has any tasks in it.%0A   */%0A
   @Override%0A
@@ -3302,24 +3302,185 @@
 k ) );%0A  %7D%0A%0A
+  /**%0A   * Remove and return the next task in queue.%0A   * This may return null if there is no tasks in the quue.%0A   *%0A   * @return the next task in queue.%0A   */%0A
   @Nullable%0A
@@ -3940,16 +3940,138 @@
 l;%0A  %7D%0A%0A
+  /**%0A   * Clear all tasks from queue and return any tasks removed.%0A   *%0A   * @return tasks removed from the queue.%0A   */%0A
   @Overr
@@ -4483,24 +4483,344 @@
 tasks;%0A  %7D%0A%0A
+  /**%0A   * Return a stream containing tasks ordered as they would be executed.%0A   * This method may be very slow and should not be invoked during production compiles.%0A   * It is only expected to be called from invariant checking code.%0A   *%0A   * @return a stream containing tasks ordered as they would be executed.%0A   */%0A
   @Nonnull%0A 

Copy javadocs to the implementation