package com.metamx.druid.client.cache;

import com.google.caliper.Param;
import com.google.caliper.Runner;
import com.google.caliper.SimpleBenchmark;
import com.google.common.collect.Lists;
import com.metamx.common.Pair;
import net.spy.memcached.AddrUtil;
import net.spy.memcached.ConnectionFactoryBuilder;
import net.spy.memcached.DefaultHashAlgorithm;
import net.spy.memcached.FailureMode;
import net.spy.memcached.MemcachedClient;
import net.spy.memcached.MemcachedClientIF;
import net.spy.memcached.transcoders.SerializingTranscoder;

import java.nio.ByteBuffer;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.TimeUnit;

public class MemcachedCacheBrokerBenchmark extends SimpleBenchmark
{
  private static final String BASE_KEY = "test_2012-11-26T00:00:00.000Z_2012-11-27T00:00:00.000Z_2012-11-27T04:11:25.979Z_";
  public static final String IDENTIFIER = "default";

  private MemcachedCacheBroker cache;
  private MemcachedClientIF client;

  private static byte[] randBytes;

  @Param({"localhost:11211"}) String hosts;

  // object size in kB
  @Param({"1", "5", "10", "40"}) int objectSize;
  @Param({"100", "1000"}) int objectCount;

  @Override
  protected void setUp() throws Exception
  {
    SerializingTranscoder transcoder = new SerializingTranscoder(
        50 * 1024 * 1024 // 50 MB
    );
    // disable compression
    transcoder.setCompressionThreshold(Integer.MAX_VALUE);

    System.out.println(String.format("Using memcached hosts [%s]", hosts));

    client = new MemcachedClient(
        new ConnectionFactoryBuilder().setProtocol(ConnectionFactoryBuilder.Protocol.BINARY)
                                      .setHashAlg(DefaultHashAlgorithm.FNV1A_64_HASH)
                                      .setLocatorType(ConnectionFactoryBuilder.Locator.CONSISTENT)
                                      .setDaemon(true)
                                      .setFailureMode(FailureMode.Retry)
                                      .setTranscoder(transcoder)
                                      .setShouldOptimize(true)
                                      .build(),
        AddrUtil.getAddresses(hosts)
    );

    cache = new MemcachedCacheBroker(
        client,
        30000, // 30 seconds
        3600 // 1 hour
    );


    randBytes = new byte[objectSize * 1024];
    new Random(0).nextBytes(randBytes);
  }

  @Override
  protected void tearDown() throws Exception
  {
    client.flush();
    client.shutdown();
  }

  public void timePutObjects(int reps) {
    for(int i = 0; i < reps; ++i) {
      for(int k = 0; k < objectCount; ++k) {
        String key = BASE_KEY + i;
        cache.put(IDENTIFIER, key.getBytes(), randBytes);
      }
      // make sure the write queue is empty
      client.waitForQueues(1, TimeUnit.HOURS);
    }
  }

  public long timeGetObject(int reps) {
    byte[] bytes = null;
    long count = 0;
    for (int i = 0; i < reps; i++) {
      for(int k = 0; k < objectCount; ++k) {
        String key = BASE_KEY + k;
        bytes = cache.get(IDENTIFIER, key.getBytes());
        count += bytes.length;
      }
    }
    return count;
  }

  public long timeBulkGetObjects(int reps) {
    long count = 0;
    for (int i = 0; i < reps; i++) {
      List<Pair<String, ByteBuffer>> keys = Lists.newArrayList();
      for(int k = 0; k < objectCount; ++k) {
        String key = BASE_KEY + k;
        keys.add(Pair.of(IDENTIFIER, ByteBuffer.wrap(key.getBytes())));
      }
      Map<Pair<String, ByteBuffer>, byte[]> results = cache.getBulk(keys);
      for(byte[] bytes : results.values()) count += bytes.length;
    }
    return count;
  }

  public static void main(String[] args) throws Exception {
    Runner.main(MemcachedCacheBrokerBenchmark.class, args);
  }
}

@@ -1445,85 +1445,8 @@
 );%0A%0A
-    System.out.println(String.format(%22Using memcached hosts %5B%25s%5D%22, hosts));%0A%0A
     
@@ -1475,16 +1475,16 @@
 Client(%0A
+
         
@@ -2383,37 +2383,36 @@
 ent.
-flush();%0A    client.shutdown(
+shutdown(1, TimeUnit.MINUTES
 );%0A 
@@ -2569,17 +2569,17 @@
 E_KEY + 
-i
+k
 ;%0A      
@@ -3478,16 +3478,63 @@
     for(
+Pair%3CString, ByteBuffer%3E key : keys) %7B%0A        
 byte%5B%5D b
@@ -3538,17 +3538,17 @@
 %5D bytes 
-:
+=
  results
@@ -3552,17 +3552,25 @@
 lts.
-values())
+get(key);%0A       
  cou
@@ -3585,24 +3585,32 @@
 tes.length;%0A
+      %7D%0A
     %7D%0A    re

fix duplicate keys, shutdown gracefully and make sure we check all multiget keys in memcached benchmark