package io.vertx.ext.mongo.impl;

import com.mongodb.async.AsyncBatchCursor;
import com.mongodb.async.SingleResultCallback;
import com.mongodb.async.client.MongoIterable;
import io.vertx.core.Context;
import io.vertx.core.Handler;
import io.vertx.core.json.JsonObject;
import io.vertx.core.streams.ReadStream;

import java.util.ArrayDeque;
import java.util.Collections;
import java.util.Deque;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

/**
 * @author Thomas Segismont
 */
class MongoIterableStream implements ReadStream<JsonObject> {

  private static final int BATCH_SIZE = 10;

  private final Context context;
  private final MongoIterable<JsonObject> mongoIterable;

  private AsyncBatchCursor<JsonObject> batchCursor;
  private Deque<JsonObject> queue;
  private Handler<JsonObject> dataHandler;
  private Handler<Throwable> exceptionHandler;
  private Handler<Void> endHandler;
  private boolean paused;
  private boolean readInProgress;
  private boolean closed;

  MongoIterableStream(Context context, MongoIterable<JsonObject> mongoIterable) {
    this.context = context;
    this.mongoIterable = mongoIterable;
  }

  @Override
  public synchronized MongoIterableStream exceptionHandler(Handler<Throwable> handler) {
    checkClosed();
    this.exceptionHandler = handler;
    return this;
  }

  private void checkClosed() {
    if (closed) {
      throw new IllegalArgumentException("Stream is closed");
    }
  }

  @Override
  public synchronized MongoIterableStream handler(Handler<JsonObject> handler) {
    checkClosed();
    if (handler == null) {
      close();
    } else {
      dataHandler = handler;
      SingleResultCallback<AsyncBatchCursor<JsonObject>> callback = (result, t) -> {
        context.runOnContext(v -> {
          synchronized (this) {
            if (t != null) {
              close();
              handleException(t);
            } else {
              batchCursor = result;
              batchCursor.setBatchSize(BATCH_SIZE);
              if (canRead()) {
                doRead();
              }
            }
          }
        });
      };
      try {
        mongoIterable.batchCursor(callback);
      } catch (Exception e) {
        close();
        handleException(e);
      }
    }
    return this;
  }

  private boolean canRead() {
    return !paused && !closed;
  }

  @Override
  public synchronized MongoIterableStream pause() {
    checkClosed();
    paused = true;
    return this;
  }

  @Override
  public synchronized MongoIterableStream resume() {
    checkClosed();
    if (paused) {
      paused = false;
      if (dataHandler != null) {
        doRead();
      }
    }
    return this;
  }

  private synchronized void doRead() {
    if (readInProgress) {
      return;
    }
    readInProgress = true;
    if (queue == null) {
      queue = new ArrayDeque<>(BATCH_SIZE);
    }
    if (!queue.isEmpty()) {
      context.runOnContext(v -> emitQueued());
      return;
    }
    context.<List<JsonObject>>executeBlocking(fut -> {
      batchCursor.next((result, t) -> {
        if (t != null) {
          fut.fail(t);
        } else {
          fut.complete(result == null ? Collections.emptyList() : result);
        }
      });
    }, false, ar -> {
      synchronized (this) {
        if (ar.succeeded()) {
          queue.addAll(ar.result());
          if (queue.isEmpty()) {
            close();
            if (endHandler != null) {
              endHandler.handle(null);
            }
          } else {
            emitQueued();
          }
        } else {
          close();
          handleException(ar.cause());
        }
      }
    });
  }

  private void handleException(Throwable cause) {
    if (exceptionHandler != null) {
      exceptionHandler.handle(cause);
    }
  }

  private synchronized void emitQueued() {
    while (!queue.isEmpty() && canRead()) {
      dataHandler.handle(queue.remove());
    }
    readInProgress = false;
    if (canRead()) {
      doRead();
    }
  }

  @Override
  public synchronized MongoIterableStream endHandler(Handler<Void> handler) {
    endHandler = handler;
    return this;
  }

  private void close() {
    closed = true;
    AtomicReference<AsyncBatchCursor> cursorRef = new AtomicReference<>();
    context.executeBlocking(fut -> {
      synchronized (this) {
        cursorRef.set(batchCursor);
      }
      AsyncBatchCursor cursor = cursorRef.get();
      if (cursor != null) {
        cursor.close();
      }
      fut.complete();
    }, false, null);
  }
}

@@ -698,16 +698,75 @@
 rable;%0A%0A
+  // All the following fields are guarded by this instance%0A
   privat
@@ -809,16 +809,16 @@
 Cursor;%0A
-
   privat
@@ -1383,32 +1383,87 @@
 turn this;%0A  %7D%0A%0A
+  // Always called from a synchronized method or block%0A
   private void c
@@ -2392,32 +2392,87 @@
 turn this;%0A  %7D%0A%0A
+  // Always called from a synchronized method or block%0A
   private boolea
@@ -2848,32 +2848,87 @@
 turn this;%0A  %7D%0A%0A
+  // Always called from a synchronized method or block%0A
   private synchr
@@ -3869,24 +3869,79 @@
    %7D);%0A  %7D%0A%0A
+  // Always called from a synchronized method or block%0A
   private vo
@@ -4059,24 +4059,79 @@
 %0A    %7D%0A  %7D%0A%0A
+  // Always called from a synchronized method or block%0A
   private sy
@@ -4467,16 +4467,71 @@
 s;%0A  %7D%0A%0A
+  // Always called from a synchronized method or block%0A
   privat

Document synchronization in MongoIterableStream