/*

  * #%L

  * BatchedAsyncMongoCollectionImpl.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.lang.reflect.InvocationHandler;

 import java.lang.reflect.Method;

 import java.lang.reflect.Proxy;

 import java.util.ArrayList;

 import java.util.Collections;

 import java.util.LinkedList;

 import java.util.List;

 import java.util.concurrent.CancellationException;

 import java.util.concurrent.Future;

 import com.allanbank.mongodb.BatchedAsyncMongoCollection;

 import com.allanbank.mongodb.Callback;

 import com.allanbank.mongodb.Durability;

 import com.allanbank.mongodb.MongoDatabase;

 import com.allanbank.mongodb.MongoDbException;

 import com.allanbank.mongodb.Version;

 import com.allanbank.mongodb.bson.Document;

 import com.allanbank.mongodb.builder.BatchedWrite;

 import com.allanbank.mongodb.builder.BatchedWrite.Bundle;

 import com.allanbank.mongodb.builder.BatchedWriteMode;

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

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

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

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

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

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

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

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

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

 /**

  * BatchedAsyncMongoCollectionImpl provides the implementation for the

  * {@link BatchedAsyncMongoCollection}.

  * 

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

  */

 public class BatchedAsyncMongoCollectionImpl extends

         AbstractAsyncMongoCollection implements BatchedAsyncMongoCollection {

     /** The interfaces to implement via the proxy. */

     private static final Class<?>[] CLIENT_INTERFACE = new Class[] { Client.class };

     /** set to true to batch deletes. */

     private boolean myBatchDeletes = false;

     /** Set to true to batch updates. */

     private boolean myBatchUpdates = false;

     /** The mode for the writes. */

     private BatchedWriteMode myMode = BatchedWriteMode.SERIALIZE_AND_CONTINUE;

     /**

      * Creates a new BatchedAsyncMongoCollectionImpl.

      * 

      * @param client

      *            The client for interacting with MongoDB.

      * @param database

      *            The database we interact with.

      * @param name

      *            The name of the collection we interact with.

      */

     public BatchedAsyncMongoCollectionImpl(final Client client,

             final MongoDatabase database, final String name) {

         super((Client) Proxy.newProxyInstance(

                 BatchedAsyncMongoCollectionImpl.class.getClassLoader(),

                 CLIENT_INTERFACE, new CaptureClientHandler(client)), database,

                 name);

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to clear any pending messages without sending them to MongoDB.

      * </p>

      */

     @Override

     public void cancel() {

         final InvocationHandler handler = Proxy.getInvocationHandler(myClient);

         if (handler instanceof CaptureClientHandler) {

             ((CaptureClientHandler) handler).clear();

         }

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to flush any pending messages to a real serialized client.

      * </p>

      */

     @Override

     public void close() throws MongoDbException {

         flush();

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to flush any pending messages to a real serialized client.

      * </p>

      */

     @Override

     public void flush() throws MongoDbException {

         final InvocationHandler handler = Proxy.getInvocationHandler(myClient);

         if (handler instanceof CaptureClientHandler) {

             ((CaptureClientHandler) handler).flush(this);

         }

     }

     /**

      * Returns the mode for the batched writes.

      * 

      * @return The mode for the batched writes.

      */

     public BatchedWriteMode getMode() {

         return myMode;

     }

     /**

      * Returns true if the deletes should be batched.

      * 

      * @return True if the deletes should be batched.

      */

     public boolean isBatchDeletes() {

         return myBatchDeletes;

     }

     /**

      * Returns true if the updates should be batched.

      * 

      * @return True if the updates should be batched.

      */

     public boolean isBatchUpdates() {

         return myBatchUpdates;

     }

     /**

      * {@inheritDoc}

      */

     @Override

     public void setBatchDeletes(final boolean batchDeletes) {

         myBatchDeletes = batchDeletes;

     }

     /**

      * {@inheritDoc}

      */

     @Override

     public void setBatchUpdates(final boolean batchUpdates) {

         myBatchUpdates = batchUpdates;

     }

     /**

      * {@inheritDoc}

      */

     @Override

     public void setMode(final BatchedWriteMode mode) {

         myMode = mode;

     }

     /**

      * {@inheritDoc}

      * <p>

      * Overridden to return false to force the {@link AbstractMongoOperations}

      * class to always use the legacy {@link Insert}, {@link Update}, and

      * {@link Delete} messages. The {@code CaptureClientHandler.optimize()} will

      * convert those operations to bulk write commands as appropriate.

      */

     @Override

     protected boolean useWriteCommand() {

         return false;

     }

     /**

      * CaptureClientHandler provides an {@link InvocationHandler} to capture all

      * send requests and defer them until flushed.

      * 

      * @copyright 2013, Allanbank Consulting, Inc., All Rights Reserved

      */

     private static class CaptureClientHandler implements InvocationHandler {

         /** The first version to support batch write commands. */

         public static final Version BATCH_WRITE_VERSION = Version

                 .parse("2.5.4");

         /** The collection we are proxying. */

         private BatchedAsyncMongoCollectionImpl myCollection;

         /** The real (e.g., user's) callbacks. */

         private List<Callback<Reply>> myRealCallbacks;

         /**

          * The {@link Client} implementation to delegate to when sending

          * messages or handling other method calls.

          */

         private final Client myRealClient;

         /** The final results from the callback. */

         private List<Object> myResults;

         /**

          * The {@link Client} implementation to delegate to when sending

          * messages or handling other method calls.

          */

         private final List<Object[]> mySendArgs;

         /** The batched writer we are building. */

         private final BatchedWrite.Builder myWrite;

         /**

          * Creates a new CaptureClientHandler.

          * 

          * @param realClient

          *            The {@link Client} implementation to delegate to when

          *            sending messages or handling other method calls.

          */

         public CaptureClientHandler(final Client realClient) {

             myRealClient = realClient;

             myRealCallbacks = null;

             myResults = null;

             mySendArgs = new LinkedList<Object[]>();

             myWrite = BatchedWrite.builder();

         }

         /**

          * Clears the pending messages without sending them to MongoDB.

          */

         public synchronized void clear() {

             final List<Object[]> copy = new ArrayList<Object[]>(mySendArgs);

             mySendArgs.clear();

             myWrite.reset();

             myResults = null;

             myRealCallbacks = null;

             myCollection = null;

             for (final Object[] args : copy) {

                 final Object lastArg = args[args.length - 1];

                 if (lastArg instanceof Future<?>) {

                     ((Future<?>) lastArg).cancel(false);

                 }

                 else if (lastArg instanceof Callback<?>) {

                     ((Callback<?>) lastArg)

                             .exception(new CancellationException(

                                     "Batch request cancelled."));

                 }

             }

         }

         /**

          * Flushes the pending messages to a serialized client.

          * 

          * @param collection

          *            The Collection the we are flushing operations for.

          */

         public synchronized void flush(

                 final BatchedAsyncMongoCollectionImpl collection) {

             // Use a serialized client to keep all of the messages on a single

             // connection as much as possible.

             SerialClientImpl serialized;

             if (myRealClient instanceof SerialClientImpl) {

                 serialized = (SerialClientImpl) myRealClient;

             }

             else {

                 serialized = new SerialClientImpl((ClientImpl) myRealClient);

             }

             try {

                 // Send the optimized requests.

                 final List<Object> optimized = optimize(collection);

                 for (final Object toSend : optimized) {

                     if (toSend instanceof BatchedWriteCallback) {

                         final BatchedWriteCallback cb = (BatchedWriteCallback) toSend;

                         cb.setClient(serialized);

                         cb.send();

                     }

                     else if (toSend instanceof Object[]) {

                         final Object[] sendArg = (Object[]) toSend;

                         if (sendArg.length == 2) {

                             serialized.send((Message) sendArg[0],

                                     (ReplyCallback) sendArg[1]);

                         }

                         else {

                             serialized.send((Message) sendArg[0],

                                     (Message) sendArg[1],

                                     (ReplyCallback) sendArg[2]);

                         }

                     }

                 }

             }

             finally {

                 clear();

             }

         }

         /**

          * {@inheritDoc}

          * <p>

          * Overridden to batch all {@link Client#send} operations.

          * </p>

          */

         @Override

         public synchronized Object invoke(final Object proxy,

                 final Method method, final Object[] args) throws Throwable {

             final String methodName = method.getName();

             if (methodName.equals("send")) {

                 mySendArgs.add(args);

                 return null;

             }

             return method.invoke(myRealClient, args);

         }

         /**

          * Adds a delete to the batch.

          * 

          * @param delete

          *            The delete to add to the batch.

          * @param args

          *            The raw send() arguments.

          */

         private void addDelete(final Delete delete, final Object[] args) {

             updateDurability(args);

             myRealCallbacks.add(extractCallback(args));

             myWrite.delete(delete.getQuery(), delete.isSingleDelete());

         }

         /**

          * Adds an insert to the batch.

          * 

          * @param insert

          *            The insert to add to the batch.

          * @param args

          *            The raw send() arguments.

          */

         private void addInsert(final Insert insert, final Object[] args) {

             updateDurability(args);

             final int docCount = insert.getDocuments().size();

             Callback<Reply> cb = extractCallback(args);

             final boolean breakBatch = (cb != null)

                     && insert.isContinueOnError() && (docCount > 1);

             if (breakBatch) {

                 closeBatch();

                 myWrite.setMode(BatchedWriteMode.SERIALIZE_AND_STOP);

             }

             else {

                 cb = new BatchedInsertCountingCallback(cb, docCount);

             }

             for (final Document doc : insert.getDocuments()) {

                 myWrite.insert(doc);

                 myRealCallbacks.add(cb);

             }

             if (breakBatch) {

                 closeBatch();

             }

         }

         /**

          * Adds an update to the batch.

          * 

          * @param update

          *            The update to add to the batch.

          * @param args

          *            The raw send() arguments.

          */

         private void addUpdate(final Update update, final Object[] args) {

             updateDurability(args);

             myRealCallbacks.add(extractCallback(args));

             myWrite.update(update.getQuery(), update.getUpdate(),

                     update.isMultiUpdate(), update.isUpsert());

         }

         /**

          * Closes the current batch of operations and re-initializes the batched

          * writer.

          */

         private void closeBatch() {

             final ClusterStats stats = myRealClient.getClusterStats();

             final BatchedWrite w = myWrite.build();

             final List<Bundle> bundles = w.toBundles(myCollection.getName(),

                     stats.getSmallestMaxBsonObjectSize(),

                     stats.getSmallestMaxBatchedWriteOperations());

             if (!bundles.isEmpty()) {

                 final BatchedWriteCallback cb = new BatchedWriteCallback(

                         myCollection.getDatabaseName(), myCollection.getName(),

                         myRealCallbacks, w, bundles);

                 myResults.add(cb);

             }

             myWrite.reset();

             myWrite.setMode(myCollection.getMode());

             myRealCallbacks.clear();

         }

         /**

          * Extracts the callback from the write arguments. If the write has a

          * {@link Callback} then it will be the last argument.

          * 

          * @param args

          *            The arguments for the original {@link Client#send} call.

          * @return The callback for the call. Returns null if there is no

          *         {@link Callback}.

          */

         private Callback<Reply> extractCallback(final Object[] args) {

             final Object cb = args[args.length - 1];

             if (cb instanceof AbstractReplyCallback<?>) {

                 return (AbstractReplyCallback<?>) args[2];

             }

             return null;

         }

         /**

          * Tries the optimize the messages we will send to the server by

          * coalescing the sequential insert, update and delete messages into the

          * batched write commands of the same name.

          * 

          * @param collection

          *            The collection we are sending requests to.

          * @return The list of optimized messages.

          */

         private List<Object> optimize(

                 final BatchedAsyncMongoCollectionImpl collection) {

             if (mySendArgs.isEmpty()) {

                 return Collections.emptyList();

             }

             final ClusterStats stats = myRealClient.getClusterStats();

             final Version minVersion = stats.getServerVersionRange()

                     .getLowerBounds();

             final boolean supportsBatch = BATCH_WRITE_VERSION

                     .compareTo(minVersion) <= 0;

             if (supportsBatch) {

                 myCollection = collection;

                 myWrite.reset();

                 myWrite.setMode(collection.getMode());

                 myResults = new ArrayList<Object>(mySendArgs.size());

                 myRealCallbacks = new ArrayList<Callback<Reply>>(

                         mySendArgs.size());

                 while (!mySendArgs.isEmpty()) {

                     final Object[] args = mySendArgs.remove(0);

                     if (args[0] instanceof Insert) {

                         addInsert((Insert) args[0], args);

                     }

                     else if (collection.isBatchUpdates()

                             && (args[0] instanceof Update)) {

                         addUpdate((Update) args[0], args);

                     }

                     else if (collection.isBatchDeletes()

                             && (args[0] instanceof Delete)) {

                         addDelete((Delete) args[0], args);

                     }

                     else {

                         closeBatch();

                         myResults.add(args);

                     }

                     if (collection.getMode() == BatchedWriteMode.SERIALIZE_AND_STOP) {

                         closeBatch();

                     }

                 }

                 closeBatch();

             }

             else {

                 myResults = new ArrayList<Object>(mySendArgs.size());

                 myResults.addAll(mySendArgs);

                 // Clear the sendArgs or they will get notified of a cancel.

                 mySendArgs.clear();

             }

             return myResults;

         }

         /**

          * Updates the durability for the batch. If the durability changes

          * mid-batch then we force a break in the batch.

          * 

          * @param args

          *            The arguments for the send() call. The

          *            {@link GetLastError} will be the second of three

          *            arguments.

          */

         private void updateDurability(final Object[] args) {

             Durability active = myWrite.getDurability();

             if ((args.length == 3) && (args[1] instanceof GetLastError)) {

                 final GetLastError error = (GetLastError) args[1];

                 final Durability d = Durability.valueOf(error.getQuery()

                         .toString());

                 if (active == null) {

                     active = d;

                     myWrite.setDurability(active);

                 }

                 else if (!d.equals(active) && !d.equals(Durability.ACK)

                         && !d.equals(Durability.NONE)) {

                     closeBatch();

                     active = d;

                     myWrite.setDurability(active);

                 }

             } // else Durability is none or not applicable.

         }

     }

 }