/*
 * 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.dse.driver.internal.core.cql.reactive;

import com.datastax.dse.driver.api.core.cql.reactive.ReactiveRow;
import com.datastax.dse.driver.internal.core.util.concurrent.BoundedConcurrentQueue;
import com.datastax.oss.driver.api.core.AsyncPagingIterable;
import com.datastax.oss.driver.api.core.cql.ColumnDefinitions;
import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
import com.datastax.oss.driver.api.core.cql.Row;
import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
import com.datastax.oss.driver.shaded.guava.common.collect.Iterators;
import edu.umd.cs.findbugs.annotations.NonNull;
import edu.umd.cs.findbugs.annotations.Nullable;
import java.util.Collections;
import java.util.Iterator;
import java.util.Objects;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import net.jcip.annotations.ThreadSafe;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A single-subscriber subscription that executes one single query and emits all the returned rows.
 *
 * <p>This class can handle both continuous and non-continuous result sets.
 */
@ThreadSafe
public class ReactiveResultSetSubscription<ResultSetT extends AsyncPagingIterable<Row, ResultSetT>>
    implements Subscription {

  private static final Logger LOG = LoggerFactory.getLogger(ReactiveResultSetSubscription.class);

  private static final int MAX_ENQUEUED_PAGES = 4;

  /** Tracks the number of items requested by the subscriber. */
  private final AtomicLong requested = new AtomicLong(0);

  /** The pages received so far, with a maximum of MAX_ENQUEUED_PAGES elements. */
  private final BoundedConcurrentQueue<Page<ResultSetT>> pages =
      new BoundedConcurrentQueue<>(MAX_ENQUEUED_PAGES);

  /**
   * Used to signal that a thread is currently draining, i.e., emitting items to the subscriber.
   * When it is zero, that means there is no ongoing emission. This mechanism effectively serializes
   * access to the drain() method, and also keeps track of missed attempts to enter it, since each
   * thread that attempts to drain will increment this counter.
   *
   * @see #drain()
   */
  private final AtomicInteger draining = new AtomicInteger(0);

  /**
   * Waited upon by the driver and completed when the subscriber requests its first item.
   *
   * <p>Used to hold off emitting results until the subscriber issues its first request for items.
   * Since this future is only completed from {@link #request(long)}, this effectively conditions
   * the enqueueing of the first page to the reception of the subscriber's first request.
   *
   * <p>This mechanism avoids sending terminal signals before a request is made when the stream is
   * empty. Note that as per 2.9, "a Subscriber MUST be prepared to receive an onComplete signal
   * with or without a preceding Subscription.request(long n) call." However, the TCK considers it
   * as unfair behavior.
   *
   * @see #start(Callable)
   * @see #request(long)
   */
  private final CompletableFuture<Void> firstSubscriberRequestArrived = new CompletableFuture<>();

  /** non-final because it has to be de-referenced, see {@link #clear()}. */
  private volatile Subscriber<? super ReactiveRow> mainSubscriber;

  private volatile Subscriber<ColumnDefinitions> columnDefinitionsSubscriber;

  private volatile Subscriber<ExecutionInfo> executionInfosSubscriber;

  private volatile Subscriber<Boolean> wasAppliedSubscriber;

  /**
   * Set to true when the subscription is cancelled, which happens when an error is encountered,
   * when the result set is fully consumed and the subscription terminates, or when the subscriber
   * manually calls {@link #cancel()}.
   */
  private volatile boolean cancelled = false;

  ReactiveResultSetSubscription(
      @NonNull Subscriber<? super ReactiveRow> mainSubscriber,
      @NonNull Subscriber<ColumnDefinitions> columnDefinitionsSubscriber,
      @NonNull Subscriber<ExecutionInfo> executionInfosSubscriber,
      @NonNull Subscriber<Boolean> wasAppliedSubscriber) {
    this.mainSubscriber = mainSubscriber;
    this.columnDefinitionsSubscriber = columnDefinitionsSubscriber;
    this.executionInfosSubscriber = executionInfosSubscriber;
    this.wasAppliedSubscriber = wasAppliedSubscriber;
  }

  /**
   * Starts the query execution.
   *
   * <p>Must be called immediately after creating the subscription, but after {@link
   * Subscriber#onSubscribe(Subscription)}.
   *
   * @param firstPage The future that, when complete, will produce the first page.
   */
  void start(@NonNull Callable<CompletionStage<ResultSetT>> firstPage) {
    firstSubscriberRequestArrived.thenAccept(
        (aVoid) -> fetchNextPageAndEnqueue(new Page<>(firstPage), true));
  }

  @Override
  public void request(long n) {
    // As per 3.6: after the Subscription is cancelled, additional
    // calls to request() MUST be NOPs.
    if (!cancelled) {
      if (n < 1) {
        // Validate request as per rule 3.9
        doOnError(
            new IllegalArgumentException(
                mainSubscriber
                    + " violated the Reactive Streams rule 3.9 by requesting a non-positive number of elements."));
      } else {
        // As per rule 3.17, when demand overflows Long.MAX_VALUE
        // it can be treated as "effectively unbounded"
        ReactiveOperators.addCap(requested, n);
        // Set the first future to true if not done yet.
        // This will make the first page of results ready for consumption,
        // see start().
        // As per 2.7 it is the subscriber's responsibility to provide
        // external synchronization when calling request(),
        // so the check-then-act idiom below is good enough
        // (and besides, complete() is idempotent).
        if (!firstSubscriberRequestArrived.isDone()) {
          firstSubscriberRequestArrived.complete(null);
        }
        drain();
      }
    }
  }

  @Override
  public void cancel() {
    // As per 3.5: Subscription.cancel() MUST respect the responsiveness of
    // its caller by returning in a timely manner, MUST be idempotent and
    // MUST be thread-safe.
    if (!cancelled) {
      cancelled = true;
      if (draining.getAndIncrement() == 0) {
        // If nobody is draining, clear now;
        // otherwise, the draining thread will notice
        // that the cancelled flag was set
        // and will clear for us.
        clear();
      }
    }
  }

  /**
   * Attempts to drain available items, i.e. emit them to the subscriber.
   *
   * <p>Access to this method is serialized by the field {@link #draining}: only one thread at a
   * time can drain, but threads that attempt to drain while other thread is already draining
   * increment that field; the draining thread, before finishing its work, checks for such failed
   * attempts and triggers another round of draining if that was the case.
   *
   * <p>The loop is interrupted when 1) the requested amount has been met or 2) when there are no
   * more items readily available or 3) the subscription has been cancelled.
   *
   * <p>The loop also checks for stream exhaustion and emits a terminal {@code onComplete} signal in
   * this case.
   *
   * <p>This method may run on a driver IO thread when invoked from {@link
   * #fetchNextPageAndEnqueue(Page, boolean)}, or on a subscriber thread, when invoked from {@link
   * #request(long)}.
   */
  @SuppressWarnings("ConditionalBreakInInfiniteLoop")
  private void drain() {
    // As per 3.4: this method SHOULD respect the responsiveness
    // of its caller by returning in a timely manner.
    // We accomplish this by a wait-free implementation.
    if (draining.getAndIncrement() != 0) {
      // Someone else is already draining, so do nothing,
      // the other thread will notice that we attempted to drain.
      // This also allows to abide by rule 3.3 and avoid
      // cycles such as request() -> onNext() -> request() etc.
      return;
    }
    int missed = 1;
    // Note: when termination is detected inside this loop,
    // we MUST call clear() manually.
    for (; ; ) {
      // The requested number of items at this point
      long r = requested.get();
      // The number of items emitted thus far
      long emitted = 0L;
      while (emitted != r) {
        if (cancelled) {
          clear();
          return;
        }
        Object result;
        try {
          result = tryNext();
        } catch (Throwable t) {
          doOnError(t);
          clear();
          return;
        }
        if (result == null) {
          break;
        }
        if (result instanceof Throwable) {
          doOnError((Throwable) result);
          clear();
          return;
        }
        doOnNext((ReactiveRow) result);
        emitted++;
      }
      if (isExhausted()) {
        doOnComplete();
        clear();
        return;
      }
      if (cancelled) {
        clear();
        return;
      }
      if (emitted != 0) {
        // if any item was emitted, adjust the requested field
        ReactiveOperators.subCap(requested, emitted);
      }
      // if another thread tried to call drain() while we were busy,
      // then we should do another drain round.
      missed = draining.addAndGet(-missed);
      if (missed == 0) {
        break;
      }
    }
  }

  /**
   * Tries to return the next item, if one is readily available, and returns {@code null} otherwise.
   *
   * <p>Cannot run concurrently due to the {@link #draining} field.
   */
  @Nullable
  private Object tryNext() {
    Page current = pages.peek();
    if (current != null) {
      if (current.hasMoreRows()) {
        return current.nextRow();
      } else if (current.hasMorePages()) {
        // Discard current page as it is consumed.
        // Don't discard the last page though as we need it
        // to test isExhausted(). It will be GC'ed when a terminal signal
        // is issued anyway, so that's no big deal.
        if (pages.poll() == null) {
          throw new AssertionError("Queue is empty, this should not happen");
        }
        // if the next page is readily available,
        // serve its first row now, no need to wait
        // for the next drain.
        return tryNext();
      }
    }
    // No items available right now.
    return null;
  }

  /**
   * Returns {@code true} when the entire stream has been consumed and no more items can be emitted.
   * When that is the case, a terminal signal is sent.
   *
   * <p>Cannot run concurrently due to the draining field.
   */
  private boolean isExhausted() {
    Page current = pages.peek();
    // Note: current can only be null when:
    // 1) we are waiting for the first page and it hasn't arrived yet;
    // 2) we just discarded the current page, but the next page hasn't arrived yet.
    // In any case, a null here means it is not the last page, since the last page
    // stays in the queue until the very end of the operation.
    return current != null && !current.hasMoreRows() && !current.hasMorePages();
  }

  /**
   * Runs on a subscriber thread initially, see {@link #start(Callable)}. Subsequent executions run
   * on the thread that completes the pair of futures [current.fetchNextPage, pages.offer] and
   * enqueues. This can be a driver IO thread or a subscriber thread; in both cases, cannot run
   * concurrently due to the fact that one can only fetch the next page when the current one is
   * arrived and enqueued.
   */
  private void fetchNextPageAndEnqueue(@NonNull Page<ResultSetT> current, boolean firstPage) {
    current
        .fetchNextPage()
        // as soon as the response arrives,
        // create the new page
        .handle(
            (rs, t) -> {
              Page<ResultSetT> page;
              if (t == null) {
                page = toPage(rs);
                executionInfosSubscriber.onNext(rs.getExecutionInfo());
                if (!page.hasMorePages()) {
                  executionInfosSubscriber.onComplete();
                }
                if (firstPage) {
                  columnDefinitionsSubscriber.onNext(rs.getColumnDefinitions());
                  columnDefinitionsSubscriber.onComplete();
                  // Avoid calling wasApplied on empty pages as some implementations may throw
                  // IllegalStateException; if the page is empty, this wasn't a CAS query, in which
                  // case, as per the method's contract, wasApplied should be true.
                  boolean wasApplied = rs.remaining() == 0 || rs.wasApplied();
                  wasAppliedSubscriber.onNext(wasApplied);
                  wasAppliedSubscriber.onComplete();
                }
              } else {
                // Unwrap CompletionExceptions created by combined futures
                if (t instanceof CompletionException) {
                  t = t.getCause();
                }
                page = toErrorPage(t);
                executionInfosSubscriber.onError(t);
                if (firstPage) {
                  columnDefinitionsSubscriber.onError(t);
                  wasAppliedSubscriber.onError(t);
                }
              }
              return page;
            })
        .thenCompose(pages::offer)
        .thenAccept(
            page -> {
              if (page.hasMorePages() && !cancelled) {
                // preemptively fetch the next page, if available
                fetchNextPageAndEnqueue(page, false);
              }
              drain();
            });
  }

  private void doOnNext(@NonNull ReactiveRow result) {
    try {
      mainSubscriber.onNext(result);
    } catch (Throwable t) {
      LOG.error(
          mainSubscriber
              + " violated the Reactive Streams rule 2.13 by throwing an exception from onNext.",
          t);
      cancel();
    }
  }

  private void doOnComplete() {
    try {
      // Then we signal onComplete as per rules 1.2 and 1.5
      mainSubscriber.onComplete();
    } catch (Throwable t) {
      LOG.error(
          mainSubscriber
              + " violated the Reactive Streams rule 2.13 by throwing an exception from onComplete.",
          t);
    }
    // We need to consider this Subscription as cancelled as per rule 1.6
    cancel();
  }

  // package-private because it can be invoked by the publisher if the subscription handshake
  // process fails.
  void doOnError(@NonNull Throwable error) {
    try {
      // Then we signal the error downstream, as per rules 1.2 and 1.4.
      mainSubscriber.onError(error);
    } catch (Throwable t) {
      t.addSuppressed(error);
      LOG.error(
          mainSubscriber
              + " violated the Reactive Streams rule 2.13 by throwing an exception from onError.",
          t);
    }
    // We need to consider this Subscription as cancelled as per rule 1.6
    cancel();
  }

  private void clear() {
    // We don't need these pages anymore and should not hold references
    // to them.
    pages.clear();
    // As per 3.13, Subscription.cancel() MUST request the Publisher to
    // eventually drop any references to the corresponding subscriber.
    // Our own publishers do not keep references to this subscription,
    // but downstream processors might do so, which is why we need to
    // defensively clear the subscriber reference when we are done.
    mainSubscriber = null;
    columnDefinitionsSubscriber = null;
    executionInfosSubscriber = null;
    wasAppliedSubscriber = null;
  }

  /**
   * Converts the received result object into a {@link Page}.
   *
   * @param rs the result object to convert.
   * @return a new page.
   */
  @NonNull
  private Page<ResultSetT> toPage(@NonNull ResultSetT rs) {
    ExecutionInfo executionInfo = rs.getExecutionInfo();
    Iterator<ReactiveRow> results =
        Iterators.transform(
            rs.currentPage().iterator(),
            row -> new DefaultReactiveRow(Objects.requireNonNull(row), executionInfo));
    return new Page<>(results, rs.hasMorePages() ? rs::fetchNextPage : null);
  }

  /** Converts the given error into a {@link Page}, containing the error as its only element. */
  @NonNull
  private Page<ResultSetT> toErrorPage(@NonNull Throwable t) {
    return new Page<>(Iterators.singletonIterator(t), null);
  }

  /**
   * A page object comprises an iterator over the page's results, and a future pointing to the next
   * page (or {@code null}, if it's the last page).
   */
  static class Page<ResultSetT extends AsyncPagingIterable<Row, ResultSetT>> {

    @NonNull final Iterator<?> iterator;

    // A pointer to the next page, or null if this is the last page.
    @Nullable final Callable<CompletionStage<ResultSetT>> nextPage;

    /** called only from start() */
    Page(@NonNull Callable<CompletionStage<ResultSetT>> nextPage) {
      this.iterator = Collections.emptyIterator();
      this.nextPage = nextPage;
    }

    Page(@NonNull Iterator<?> iterator, @Nullable Callable<CompletionStage<ResultSetT>> nextPage) {
      this.iterator = iterator;
      this.nextPage = nextPage;
    }

    boolean hasMorePages() {
      return nextPage != null;
    }

    @NonNull
    CompletionStage<ResultSetT> fetchNextPage() {
      try {
        return Objects.requireNonNull(nextPage).call();
      } catch (Exception e) {
        // This is a synchronous failure in the driver.
        // It can happen in rare cases when the driver throws an exception instead of returning a
        // failed future; e.g. if someone tries to execute a continuous paging request but the
        // protocol version in use does not support it.
        // We treat it as a failed future.
        return CompletableFutures.failedFuture(e);
      }
    }

    boolean hasMoreRows() {
      return iterator.hasNext();
    }

    @NonNull
    Object nextRow() {
      return iterator.next();
    }
  }
}
