import com.mongodb.BasicDBObject;
import com.mongodb.CommandResult;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.Mongo;
import com.mongodb.MongoOptions;
import com.mongodb.MongoURI;
import com.mongodb.WriteConcern;

import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;

public class MongoBenchTest {

    private static List<DBCollection> createCollectionList(Mongo mongo, int noOfDBs, int noOfColls) throws UnknownHostException {
        List<DBCollection> collectionList = new ArrayList<DBCollection>();
        for (int i = 0; i < noOfDBs; i++) {
            DB mDB = mongo.getDB("mongoBenchDB" + i);
            for (int j = 0; j < noOfColls; j++) {
                DBCollection mColl = mDB.getCollection("mongoBenchColl" + j);
                mColl.drop();
                mColl.ensureIndex(new BasicDBObject("randomNo", 1));
                mColl.ensureIndex(new BasicDBObject("title", 1));
                mColl.ensureIndex(new BasicDBObject("stringList", 1));
                collectionList.add(mColl);
            }
        }

        return collectionList;
    }

    private static void resetCollection(DBCollection c) {
        c.drop();
        c.ensureIndex(new BasicDBObject("randomNo", 1));
    }

    public static void main(String[] args) {

        try {
            int documentSizeFactor = 4000;  // ignored currently
            int noOfDBs = 20;
            int noOfColls = 1;
            int noOfOperations = 4000;
            int batchSize = 4000;
            int startNumThreads = 1;
            int endNumThreads = 10;
            int threadNumIncrement = 1;
            final int numRuns = 1;
            boolean writeTest = true;

            String uri = "mongodb://localhost:30000";
            if (args.length > 0) {
                uri = args[0];
            }

            uri += "/?maxPoolSize=" + endNumThreads;


            MongoOptions mo = new MongoOptions();
            mo.connectionsPerHost = endNumThreads;
            Mongo mongo = new Mongo(new MongoURI(uri));

            List<DBCollection> collectionList = createCollectionList(mongo, noOfDBs, noOfColls);

            CommandResult serverStatus = mongo.getDB("admin").command("serverStatus");

            System.out.println("Mongo server: " + uri);
            System.out.println("MongoDB server version: " + serverStatus.get("version"));
            System.out.println("Number Of threads = " + endNumThreads);
            System.out.println("Number Of operations = " + noOfOperations);
            System.out.println("Batch size = " + batchSize);
            System.out.println("Number of runs = " + numRuns);
            System.out.println();


            if (writeTest) {
                // warm up
                doWriteTest(5, documentSizeFactor, endNumThreads, 5, collectionList);

                for (int threadCount = startNumThreads; threadCount <= endNumThreads; threadCount+=threadNumIncrement) {
                    System.out.println("Thread count: " + threadCount);
                    long totalElapsedTime = 0;
                    for (int i = 1; i <= numRuns; i++) {
                        long elapsedTime = doWriteTest(noOfOperations, documentSizeFactor, threadCount, batchSize, collectionList);
                        totalElapsedTime += elapsedTime;

//                        System.out.println(" " + elapsedTime + " ms");
                    }
//                    System.out.println("Total: " + totalElapsedTime + " ms");
                    System.out.println(" Average: " + totalElapsedTime / numRuns + " ms");
                    System.out.println();
                }

                System.out.println();
                System.out.println("Done");
            }
        } catch (Exception e) {
            System.out.println("Exception occurred " + e.getMessage());
            e.printStackTrace();
        }

    }

    private static long doWriteTest(final int noOfOperations, final int documentSizeFactor, final int noOfThreads,
                                    final int batchSize,
                                    List<DBCollection> collectionList) throws InterruptedException {
        CountDownLatch ready = new CountDownLatch(noOfThreads);
        CountDownLatch go = new CountDownLatch(1);

        List<Thread> threads = new ArrayList<Thread>();
        for (int i = 0; i < noOfThreads; i++) {
            resetCollection(collectionList.get(i));
            MongoWriteThread mwt = new MongoWriteThread(collectionList.get(i), documentSizeFactor, noOfOperations, batchSize, ready, go);
            threads.add(mwt);
            mwt.start();
        }

        ready.await();

        long startTime = System.currentTimeMillis();

        go.countDown();

        for (Thread t : threads) {
            t.join();
        }

        long endTime = System.currentTimeMillis();

        return (endTime - startTime);
    }

    static class MongoWriteThread extends Thread {

        public MongoWriteThread(DBCollection coll, int docSizeFactor, int noOfOperations, int batchSize, final CountDownLatch ready, final CountDownLatch go) {
            this.coll = coll;
            this.docSizeFactor = docSizeFactor;
            this.batchSize = batchSize;
            this.ready = ready;
            this.go = go;
            this.noOfOperations = noOfOperations;
        }

        private DBCollection coll;
        // Use docSizeFactor as an input to determine document size returned by getObjectOFSize
        private int docSizeFactor;
        private final int batchSize;
        private final CountDownLatch ready;
        private final CountDownLatch go;
        private int noOfOperations;

        private long randomNo;

        private String title = "sample title";
        private String[] stringList = {"listString1", "listString2", "listString3", "listString4", "listString5", "listString6", "listString7", "listString8"};
        private String longString = "efwwefwerferwjkhwfgkrwjehfkjrhwbfjherbwfkhjerwkjhbfkejrhbwfkjrhwebfjkhrwbfkjhrwbfjkhrebwfjkhberwfjhbwrefjhberwfbherwkjfbwkjerb" +
                "klwrehflkerwhfleruwhflkwhrejfklherwldjkfhwllllerfkljerbwfljkberfklberwliugirhlubflerbwki3475gferuwbfkwjbulferhlkfhklrewhrwkleflgkehrlkghwlerhgkerflkjhrew" +
                "fluiwerhfuhgerfkuygwreyukfgekrwyugfuiqwghfiurwgbfkeuwygi7348gfbkjwbfrkjrwbhfkuyfeu4kfgku34wyfrhfglurehligkhrelwhgkerwghfkrewgfkuerwgfkerwfgkevkjqwflqrwbf" +
                "lfyurwegflrgfliruewhflrewhfleriwhflirewufhlierwflierwgfhlierwhuflriwefhlierwufhlirewuhflierwhfuliwrefhleirwuhflirehglwghletrkhgltieghelgherlwghlewihglerh" +
                "erfkljrbflkwrrlkwhelfkuhewrlkuhlkgreuhglkhewrlkuhglkrehglkrhueglkuehrwlgkhglkwerhfgliuerwhliwhilhluerbwlkgbelwrkbvlkurehwfligwerhfkygjbflkrbflkrwebflrkbge";


        private DBObject getDBObjectOfSize(long docSizeFactor) {
            DBObject dbo = new BasicDBObject();
            dbo.put("randomNo", randomNo++);
            dbo.put("title", title);
            dbo.put("stringList", stringList);
            dbo.put(longString, longString);
            return dbo;
        }

        private DBObject[] getDBObjectArray() {

            List<DBObject> dbObjectList = new ArrayList<DBObject>();

            for (int i = 0; i < noOfOperations; i++) {
                DBObject dbo = getDBObjectOfSize(docSizeFactor);
                dbObjectList.add(dbo);
            }

            return dbObjectList.toArray(new DBObject[noOfOperations]);
        }

        public void run() {
            try {
                DBObject[] dboArray = getDBObjectArray();

                ready.countDown();
                go.await();

                DBObject[] batch = new DBObject[batchSize];
                for (int batchNum = 0; batchNum < noOfOperations / batchSize; batchNum++) {
                    System.arraycopy(dboArray, batchSize * batchNum, batch, 0, batchSize);
                    coll.insert(batch, WriteConcern.SAFE);
                }

            } catch (Exception e) {
                e.printStackTrace();
            }

        }
    }
}