/*
 * Copyright DataStax, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.datastax.oss.driver.internal.core.cql;

import com.datastax.oss.driver.api.core.PagingIterable;
import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.util.Objects;
import java.util.Spliterator;
import java.util.Spliterators;
import java.util.function.Consumer;
import net.jcip.annotations.NotThreadSafe;

A Spliterator for PagingIterable instances that splits the stream in chunks of equal size. 
Type parameters:
<ElementT> – The element type of the underlying stream./**
 * A Spliterator for {@link PagingIterable} instances that splits the stream in chunks of equal
 * size.
 *
 * @param <ElementT> The element type of the underlying stream.
 */
@NotThreadSafe
public class PagingIterableSpliterator<ElementT> implements Spliterator<ElementT> {

  @NonNull
  public static <ElementT> Builder<ElementT> builder(@NonNull PagingIterable<ElementT> iterable) {
    return new Builder<>(iterable);
  }

  
The default chunk size for PagingIterableSpliterator. /** The default chunk size for {@link PagingIterableSpliterator}. */
  public static final int DEFAULT_CHUNK_SIZE = 128;

  private final PagingIterable<ElementT> iterable;
  private long estimatedSize;
  private final int chunkSize;
  private final int characteristics;

  
Creates a new PagingIterableSpliterator for the given iterable, with unknown size and default chunk size (128). 
Params:
iterable – The PagingIterable to create a spliterator for./**
   * Creates a new {@link PagingIterableSpliterator} for the given iterable, with unknown size and
   * default chunk size ({@value #DEFAULT_CHUNK_SIZE}).
   *
   * @param iterable The {@link PagingIterable} to create a spliterator for.
   */
  public PagingIterableSpliterator(@NonNull PagingIterable<ElementT> iterable) {
    this(iterable, Long.MAX_VALUE, DEFAULT_CHUNK_SIZE);
  }

  private PagingIterableSpliterator(
      @NonNull PagingIterable<ElementT> iterable, long estimatedSize, int chunkSize) {
    this.iterable = Objects.requireNonNull(iterable, "iterable cannot be null");
    this.estimatedSize = estimatedSize;
    this.chunkSize = chunkSize;
    if (estimatedSize < Long.MAX_VALUE) {
      characteristics =
          Spliterator.ORDERED
              | Spliterator.IMMUTABLE
              | Spliterator.NONNULL
              | Spliterator.SIZED
              | Spliterator.SUBSIZED;
    } else {
      characteristics = Spliterator.ORDERED | Spliterator.IMMUTABLE | Spliterator.NONNULL;
    }
  }

  @Override
  public boolean tryAdvance(Consumer<? super ElementT> action) {
    Objects.requireNonNull(action, "action cannot be null");
    ElementT row = iterable.one();
    if (row == null) {
      return false;
    }
    action.accept(row);
    return true;
  }

  @Override
  @Nullable
  public Spliterator<ElementT> trySplit() {
    if (estimatedSize != Long.MAX_VALUE && estimatedSize <= chunkSize) {
      // There is no point in splitting if the number of remaining elements is below the chunk size
      return null;
    }
    ElementT row = iterable.one();
    if (row == null) {
      return null;
    }
    Object[] array = new Object[chunkSize];
    int i = 0;
    do {
      array[i++] = row;
      if (i < chunkSize) {
        row = iterable.one();
      } else {
        break;
      }
    } while (row != null);
    if (estimatedSize != Long.MAX_VALUE) {
      estimatedSize -= i;
    }
    // Splits will also report SIZED and SUBSIZED as well.
    return Spliterators.spliterator(array, 0, i, characteristics());
  }

  @Override
  public void forEachRemaining(Consumer<? super ElementT> action) {
    iterable.iterator().forEachRemaining(action);
  }

  @Override
  public long estimateSize() {
    return estimatedSize;
  }

  @Override
  public int characteristics() {
    return characteristics;
  }

  public static class Builder<ElementT> {

    private final PagingIterable<ElementT> iterable;
    private long estimatedSize = Long.MAX_VALUE;
    private int chunkSize = DEFAULT_CHUNK_SIZE;

    Builder(@NonNull PagingIterable<ElementT> iterable) {
      this.iterable = iterable;
    }

    @NonNull
    public Builder<ElementT> withEstimatedSize(long estimatedSize) {
      Preconditions.checkArgument(estimatedSize >= 0, "estimatedSize must be >= 0");
      this.estimatedSize = estimatedSize;
      return this;
    }

    @NonNull
    public Builder<ElementT> withChunkSize(int chunkSize) {
      Preconditions.checkArgument(chunkSize > 0, "chunkSize must be > 0");
      this.chunkSize = chunkSize;
      return this;
    }

    @NonNull
    public PagingIterableSpliterator<ElementT> build() {
      return new PagingIterableSpliterator<>(iterable, estimatedSize, chunkSize);
    }
  }
}