/*

  * #%L

  * MongoIteratorImpl.java - mongodb-async-driver - Allanbank Consulting, Inc.

  * %%

  * Copyright (C) 2011 - 2014 Allanbank Consulting, 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.

  * #L%

  */

 package com.allanbank.mongodb.client;

 import java.util.ArrayList;

 import java.util.Iterator;

 import java.util.List;

 import java.util.NoSuchElementException;

 import java.util.concurrent.ExecutionException;

 import java.util.concurrent.Future;

 import com.allanbank.mongodb.MongoClient;

 import com.allanbank.mongodb.MongoDbException;

 import com.allanbank.mongodb.MongoIterator;

 import com.allanbank.mongodb.ReadPreference;

 import com.allanbank.mongodb.bson.Document;

 import com.allanbank.mongodb.bson.NumericElement;

 import com.allanbank.mongodb.bson.builder.BuilderFactory;

 import com.allanbank.mongodb.bson.builder.DocumentBuilder;

 import com.allanbank.mongodb.bson.element.StringElement;

 import com.allanbank.mongodb.client.callback.FutureReplyCallback;

 import com.allanbank.mongodb.client.message.CursorableMessage;

 import com.allanbank.mongodb.client.message.GetMore;

 import com.allanbank.mongodb.client.message.KillCursors;

 import com.allanbank.mongodb.client.message.Reply;

 import com.allanbank.mongodb.error.CursorNotFoundException;

 import com.allanbank.mongodb.util.log.Log;

 import com.allanbank.mongodb.util.log.LogFactory;

 /**

  * Iterator over the results of the MongoDB cursor.

  * 

  * @api.no This class is <b>NOT</b> part of the drivers API. This class may be

  *         mutated in incompatible ways between any two releases of the driver.

  * @copyright 2011-2014, Allanbank Consulting, Inc., All Rights Reserved

  */

 public class MongoIteratorImpl implements MongoIterator<Document> {

     /** The log for the iterator. */

     private static final Log LOG = LogFactory.getLog(MongoIteratorImpl.class);

     /** The size of batches that are requested from the servers. */

     private int myBatchSize = 0;

     /** The client for sending get_more requests to the server. */

     private final Client myClient;

     /** The name of the collection the query was originally created on. */

     private final String myCollectionName;

     /** The iterator over the current set of documents. */

     private Iterator<Document> myCurrentIterator;

     /** The original query. */

     private long myCursorId = 0;

     /** The name of the database the query was originally created on. */

     private final String myDatabaseName;

     /**

      * The maximum number of document to return from the cursor. Zero or

      * negative means all.

      */

     private int myLimit = 0;

     /** The {@link Future} that will be updated with the next set of results. */

     private FutureReplyCallback myNextReply;

     /** The read preference to subsequent requests. */

     private final ReadPreference myReadPerference;

     /**

      * Flag to shutdown this iterator gracefully without closing the cursor on

      * the server.

      */

     private boolean myShutdown = false;

     /**

      * Create a new MongoDBInterator.

      * 

      * @param originalQuery

      *            The original query being iterated over.

      * @param client

      *            The client for issuing more requests.

      * @param server

      *            The server that received the original query request.

      * @param reply

      *            The initial results of the query that are available.

      */

     public MongoIteratorImpl(final CursorableMessage originalQuery,

             final Client client, final String server, final Reply reply) {

         myNextReply = new FutureReplyCallback();

         myNextReply.callback(reply);

         myReadPerference = ReadPreference.server(server);

         myCursorId = 0;

         myClient = client;

         myCurrentIterator = null;

         myBatchSize = originalQuery.getBatchSize();

         myLimit = originalQuery.getLimit();

         myDatabaseName = originalQuery.getDatabaseName();

         myCollectionName = originalQuery.getCollectionName();

     }

     /**

      * Create a new MongoIteratorImpl from a cursor document.

      * 

      * @param client

      *            The client interface to the server.

      * @param cursorDocument

      *            The original query.

      * 

      * @see MongoIteratorImpl#asDocument()

      */

     public MongoIteratorImpl(final Document cursorDocument, final Client client) {

         final String ns = cursorDocument.get(StringElement.class,

                 NAME_SPACE_FIELD).getValue();

         String db = ns;

         String collection = ns;

         final int index = ns.indexOf('.');

         if (0 < index) {

             db = ns.substring(0, index);

             collection = ns.substring(index + 1);

         }

         myClient = client;

         myDatabaseName = db;

         myCollectionName = collection;

         myCursorId = cursorDocument.get(NumericElement.class, CURSOR_ID_FIELD)

                 .getLongValue();

         myLimit = cursorDocument.get(NumericElement.class, LIMIT_FIELD)

                 .getIntValue();

         myBatchSize = cursorDocument

                 .get(NumericElement.class, BATCH_SIZE_FIELD).getIntValue();

         myReadPerference = ReadPreference.server(cursorDocument.get(

                 StringElement.class, SERVER_FIELD).getValue());

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return the active cursor in the defined format.

      * </p>

      * 

      * @see ClientImpl#isCursorDocument(Document)

      */

     @Override

     public Document asDocument() {

         long cursorId = myCursorId;

         final Future<Reply> replyFuture = myNextReply;

         cursorId = retreiveCursorIdFromPendingRequest(cursorId, replyFuture);

         if (cursorId != 0) {

             final DocumentBuilder b = BuilderFactory.start();

             b.add(NAME_SPACE_FIELD, myDatabaseName + "." + myCollectionName);

             b.add(CURSOR_ID_FIELD, cursorId);

             b.add(SERVER_FIELD, myReadPerference.getServer());

             b.add(LIMIT_FIELD, myLimit);

             b.add(BATCH_SIZE_FIELD, myBatchSize);

             return b.build();

         }

         return null;

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to close the iterator and send a {@link KillCursors} for the

      * open cursor, if any.

      * </p>

      */

     @Override

     public void close() {

         long cursorId = myCursorId;

         final Future<Reply> replyFuture = myNextReply;

         myCurrentIterator = null;

         myNextReply = null;

         myCursorId = 0;

         cursorId = retreiveCursorIdFromPendingRequest(cursorId, replyFuture);

         if ((cursorId != 0) && !myShutdown) {

             // The user asked us to leave the cursor be.

             myClient.send(new KillCursors(new long[] { cursorId },

                     myReadPerference), null);

         }

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to get the batch size from the original query or set

      * explicitly.

      * </p>

      */

     @Override

     public int getBatchSize() {

         return myBatchSize;

     }

     /**

      * Returns the iterator's read preference which points to the original

      * server performing the query.

      * 

      * @return The iterator's read preference which points to the original

      *         server performing the query.

      */

     public ReadPreference getReadPerference() {

         return myReadPerference;

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return true if there are more documents.

      * </p>

      */

     @Override

     public boolean hasNext() {

         if (myCurrentIterator == null) {

             loadDocuments();

         }

         else if (!myCurrentIterator.hasNext() && (myNextReply != null)) {

             loadDocuments();

         }

         return myCurrentIterator.hasNext();

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return this iterator.

      * </p>

      */

     @Override

     public Iterator<Document> iterator() {

         return this;

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return the next document from the query.

      * </p>

      * 

      * @see java.util.Iterator#next()

      */

     @Override

     public Document next() {

         if (hasNext()) {

             return myCurrentIterator.next();

         }

         throw new NoSuchElementException("No more documents.");

     }

     /**

      * Computes the size for the next batch of documents to get.

      * 

      * @return The returnNex

      */

     public int nextBatchSize() {

         if ((0 < myLimit) && (myLimit <= myBatchSize)) {

             return myLimit;

         }

         return myBatchSize;

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to throw and {@link UnsupportedOperationException}.

      * </p>

      * 

      * @see java.util.Iterator#remove()

      */

     @Override

     public void remove() {

         throw new UnsupportedOperationException(

                 "Cannot remove a document via a MongoDB iterator.");

     }

     /**

      * Restarts the iterator by sending a request for more documents.

      * 

      * @throws MongoDbException

      *             On a failure to send the request for more document.

      */

     public void restart() throws MongoDbException {

         sendRequest();

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to set the batch size.

      * </p>

      */

     @Override

     public void setBatchSize(final int batchSize) {

         myBatchSize = batchSize;

     }

     /**

      * Stops the iterator after consuming any received and/or requested batches.

      * <p>

      * <b>WARNING</b>: This will leave the cursor open on the server. Users

      * should persist the state of the cursor as returned from

      * {@link #asDocument()} and restart the cursor using one of the

      * {@link MongoClient#restart(com.allanbank.mongodb.bson.DocumentAssignable)}

      * or

      * {@link MongoClient#restart(com.allanbank.mongodb.StreamCallback, com.allanbank.mongodb.bson.DocumentAssignable)}

      * methods. Use with extreme caution.

      * </p>

      * <p>

      * The iterator will naturally stop ({@link #hasNext()} will return false)

      * when the current batch and any already requested batches are finished.

      * </p>

      */

     @Override

     public void stop() {

         myShutdown = true;

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return the remaining elements as a array.

      * </p>

      */

     @Override

     public Object[] toArray() {

         final List<Document> remaining = toList();

         return remaining.toArray();

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return the remaining elements as a array.

      * </p>

      */

     @Override

     public <S> S[] toArray(final S[] to) {

         final List<Document> remaining = toList();

         return remaining.toArray(to);

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return the remaining elements as a list.

      * </p>

      */

     @Override

     public List<Document> toList() {

         final List<Document> remaining = new ArrayList<Document>();

         while (hasNext()) {

             remaining.add(next());

         }

         return remaining;

     }

     /**

      * Returns the client value.

      * 

      * @return The client value.

      */

     protected Client getClient() {

         return myClient;

     }

     /**

      * Returns the collection name.

      * 

      * @return The collection name.

      */

     protected String getCollectionName() {

         return myCollectionName;

     }

     /**

      * Returns the cursor Id value.

      * 

      * @return The cursor Id value.

      */

     protected long getCursorId() {

         return myCursorId;

     }

     /**

      * Returns the database name value.

      * 

      * @return The database name value.

      */

     protected String getDatabaseName() {

         return myDatabaseName;

     }

     /**

      * Returns the limit value.

      * 

      * @return The limit value.

      */

     protected int getLimit() {

         return myLimit;

     }

     /**

      * Loads more documents into the iterator. This iterator issues a get_more

      * command as soon as the previous results start to be used.

      * 

      * @throws RuntimeException

      *             On a failure to load documents.

      */

     protected void loadDocuments() throws RuntimeException {

         loadDocuments(true);

     }

     /**

      * Loads more documents into the iterator. This iterator issues a get_more

      * command as soon as the previous results start to be used.

      * 

      * @param blockForTailable

      *            If true then the method will recursively call itself on a

      *            tailable cursor with no results. This makes the call blocking.

      *            It false then the call will not block. This is used by the

      *            method to ensure that the outermost load blocks but the

      *            recursion is not inifinite.

      * @return The list of loaded documents.

      * 

      * @throws RuntimeException

      *             On a failure to load documents.

      */

     protected List<Document> loadDocuments(final boolean blockForTailable)

             throws RuntimeException {

         List<Document> docs;

         try {

             // Pull the reply from the future. Hopefully it is already there!

             final Reply reply = myNextReply.get();

             if (reply.isCursorNotFound() || reply.isQueryFailed()) {

                 final long cursorid = myCursorId;

                 myCursorId = 0;

                 throw new CursorNotFoundException(reply, "Cursor id ("

                         + cursorid + ") not found by the MongoDB server.");

             }

             myCursorId = reply.getCursorId();

             // Setup and iterator over the documents and adjust the limit

             // for the documents we have. Do this before the fetch again

             // so the nextBatchSize() has the updated limit.

             docs = reply.getResults();

             myCurrentIterator = docs.iterator();

             if (0 < myLimit) {

                 // Check if we have too many docs.

                 if (myLimit <= docs.size()) {

                     myCurrentIterator = docs.subList(0, myLimit).iterator();

                     if (myCursorId != 0) {

                         // Kill the cursor.

                         myClient.send(new KillCursors(

                                 new long[] { myCursorId }, myReadPerference),

                                 null);

                         myCursorId = 0;

                     }

                 }

                 myLimit -= docs.size();

             }

             // Pre-fetch the next set of documents while we iterate over the

             // documents we just got.

             if ((myCursorId != 0) && !myShutdown) {

                 sendRequest();

                 // Include the (myNextReply != null) to catch failures on the

                 // server.

                 while (docs.isEmpty() && blockForTailable

                         && (myNextReply != null)) {

                     // Tailable - Wait for a reply with documents.

                     docs = loadDocuments(false);

                 }

             }

             else {

                 // Exhausted the cursor or are shutting down - no more results.

                 myNextReply = null;

                 // Don't need to kill the cursor since we exhausted it or are

                 // shutting down.

             }

         }

         catch (final InterruptedException e) {

             throw new RuntimeException(e);

         }

         catch (final ExecutionException e) {

             throw new RuntimeException(e);

         }

         return docs;

     }

     /**

      * If the current cursor id is zero then waits for the response from the

      * pending request to determine the real cursor id.

      * 

      * @param cursorId

      *            The presumed cursor id.

      * @param replyFuture

      *            The pending reply's future.

      * @return The best known cursor id.

      */

     protected long retreiveCursorIdFromPendingRequest(final long cursorId,

             final Future<Reply> replyFuture) {

         // May not have processed any of the results yet...

         if ((cursorId == 0) && (replyFuture != null)) {

             try {

                 final Reply reply = replyFuture.get();

                 return reply.getCursorId();

             }

             catch (final InterruptedException e) {

                 LOG.warn(e, "Interrupted waiting for a query reply: {}",

                         e.getMessage());

             }

             catch (final ExecutionException e) {

                 LOG.warn(e, "Interrupted waiting for a query reply: {}",

                         e.getMessage());

             }

         }

         return cursorId;

     }

     /**

      * Sends a request for more documents.

      * 

      * @throws MongoDbException

      *             On a failure to send the request for more document.

      */

     protected void sendRequest() throws MongoDbException {

         final GetMore getMore = new GetMore(myDatabaseName, myCollectionName,

                 myCursorId, nextBatchSize(), myReadPerference);

         myNextReply = new FutureReplyCallback();

         myClient.send(getMore, myNextReply);

     }

 }