package io.vertx.ext.asyncsql.impl;

import com.github.mauricio.async.db.QueryResult;
import com.github.mauricio.async.db.ResultSet;
import com.github.mauricio.async.db.RowData;
import io.vertx.core.AsyncResult;
import io.vertx.core.Future;
import io.vertx.core.Handler;
import io.vertx.core.json.JsonArray;
import io.vertx.ext.sql.SQLRowStream;
import scala.Option;
import scala.collection.Iterator;

import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

class AsyncSQLRowStream implements SQLRowStream {

  private final ResultSet rs;
  private final Iterator<RowData> cursor;
  private List<String> columns;

  private final AtomicBoolean paused = new AtomicBoolean(false);
  private final AtomicBoolean ended = new AtomicBoolean(false);
  private final AtomicBoolean rsClosed = new AtomicBoolean(false);

  private Handler<JsonArray> handler;
  private Handler<Void> endHandler;
  private Handler<Void> rsClosedHandler;

  AsyncSQLRowStream(QueryResult qr) {
    final Option<ResultSet> rows = qr.rows();
    if (rows.isDefined()) {
      rs = rows.get();
      cursor = rs.iterator();
    } else {
      rs = null;
      cursor = null;
    }

    paused.set(true);
  }

  @Override
  public int column(String name) {
    if (rs == null) {
      throw new IndexOutOfBoundsException("'" + name + "' not found");
    }

    if (columns == null) {
      columns = ScalaUtils.toJavaList(rs.columnNames().toList());
    }
    return columns.indexOf(name);
  }

  @Override
  public SQLRowStream exceptionHandler(Handler<Throwable> handler) {
    return this;
  }

  @Override
  public SQLRowStream handler(Handler<JsonArray> handler) {
    this.handler = handler;
    // start pumping data once the handler is set
    resume();
    return this;
  }

  @Override
  public SQLRowStream pause() {
    paused.compareAndSet(false, true);
    return this;
  }

  @Override
  public SQLRowStream resume() {
    if (paused.compareAndSet(true, false)) {
      nextRow();
    }
    return this;
  }

  private void nextRow() {
    if (!paused.get()) {
      if (cursor.hasNext()) {
        handler.handle(ScalaUtils.rowToJsonArray(cursor.next()));
        nextRow();
      } else {
        // mark as ended if the handler was registered too late
        ended.set(true);
        // automatically close resources
        if (rsClosedHandler != null) {
          // only notify (since the rs is closed by the underlying driver)
          rsClosedHandler.handle(null);
        } else {
          // default behavior notify that everything ended
          close(c -> {
            if (endHandler != null) {
              endHandler.handle(null);
            }
          });
        }
      }
    }
  }

  @Override
  public SQLRowStream endHandler(Handler<Void> handler) {
    this.endHandler = handler;
    // registration was late but we're already ended, notify
    if (ended.compareAndSet(true, false)) {
      // only notify once
      endHandler.handle(null);
    }
    return this;
  }

  @Override
  public SQLRowStream resultSetClosedHandler(Handler<Void> handler) {
    this.rsClosedHandler = handler;
    // registration was late but we're already ended, notify
    if (rsClosed.compareAndSet(true, false)) {
      // only notify once
      rsClosedHandler.handle(null);
    }
    return this;
  }

  @Override
  public void moreResults() {
    // NO-OP since the underlying driver only returns 1 ResultSet
  }

  @Override
  public void close() {
    close(null);
  }

  @Override
  public void close(Handler<AsyncResult<Void>> handler) {
    // make sure we stop pumping data
    pause();
    // call the provided handler
    if (handler != null) {
      handler.handle(Future.succeededFuture());
    }
  }
}

@@ -2269,24 +2269,52 @@
 .set(true);%0A
+        rsClosed.set(true);%0A
         // a

make sure than late subscription of rs closed are also notified