package com.proofpoint.cassandra;

import com.google.common.base.Charsets;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.io.Files;
import com.google.common.net.InetAddresses;
import com.proofpoint.experimental.units.DataSize;
import com.proofpoint.node.NodeInfo;
import org.apache.cassandra.config.ConfigurationException;
import org.apache.cassandra.service.CassandraDaemon;
import org.apache.thrift.transport.TTransportException;
import org.yaml.snakeyaml.Yaml;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import javax.inject.Inject;
import java.io.File;
import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import static java.lang.String.format;
import static java.util.Arrays.asList;

public class EmbeddedCassandraServer
{
    private final CassandraDaemon cassandra;
    private final InetAddress rpcAddress;
    private final int rpcPort;

    @Inject
    public EmbeddedCassandraServer(CassandraServerConfig config, NodeInfo nodeInfo)
            throws TTransportException, IOException, InterruptedException, ConfigurationException
    {
        File directory = config.getDirectory();

        if (!directory.mkdirs() && !directory.exists()) {
            throw new IllegalStateException(format("Directory %s does not exist and cannot be created", directory));
        }

        rpcAddress = nodeInfo.getBindIp();
        rpcPort = config.getRpcPort();

        List<ImmutableMap<String, Object>> seedProvider = ImmutableList.of(ImmutableMap.<String, Object>of(
                "class_name", "org.apache.cassandra.locator.SimpleSeedProvider",
                "parameters", ImmutableList.of(ImmutableMap.of("seeds", config.getSeeds()))));

        Map<String, Object> map = ImmutableMap.<String, Object>builder()
                .put("cluster_name", config.getClusterName())
                .put("auto_bootstrap", "false")
                .put("hinted_handoff_enabled", "true")
                .put("partitioner", config.getPartitioner().getClazz().getName())
                .put("data_file_directories", asList(new File(directory, "data").getAbsolutePath()))
                .put("commitlog_directory", new File(directory, "commitlog").getAbsolutePath())
                .put("saved_caches_directory", new File(directory, "saved_caches").getAbsolutePath())
                .put("commitlog_sync", "periodic") // TODO: make configurable
                .put("commitlog_sync_period_in_ms", "10000") // TODO: make configurable
                .put("disk_access_mode", "auto")
                .put("storage_port", config.getStoragePort())
                .put("listen_address", InetAddresses.toUriString(nodeInfo.getPublicIp()))
                .put("rpc_address", InetAddresses.toUriString(rpcAddress))
                .put("rpc_port", rpcPort)
                .put("endpoint_snitch", "org.apache.cassandra.locator.SimpleSnitch") // TODO: make configurable
                .put("request_scheduler", "org.apache.cassandra.scheduler.NoScheduler")
                .put("in_memory_compaction_limit_in_mb", (int) config.getInMemoryCompactionLimit().getValue(DataSize.Unit.MEGABYTE))
                .put("sliced_buffer_size_in_kb", 64)
                .put("thrift_framed_transport_size_in_mb", 3)
                .put("thrift_max_message_length_in_mb", 4)
                .put("column_index_size_in_kb", (int) config.getColumnIndexSize().getValue(DataSize.Unit.KILOBYTE))
                .put("seed_provider", seedProvider)
                .put("encryption_options", ImmutableMap.of("internode_encryption", "none"))
                .put("rpc_timeout_in_ms", (int) config.getRpcTimeout().toMillis())
                .put("memtable_total_space_in_mb", (int) config.getMemtableTotalSpace().getValue(DataSize.Unit.MEGABYTE))
                .build();

        File configFile = new File(directory, "config.yaml");

        Files.write(new Yaml().dump(map), configFile, Charsets.UTF_8);
        System.setProperty("cassandra.config", configFile.toURI().toString());

        cassandra = new EmbeddedCassandraDaemon();
        cassandra.init(null);
    }

    @PostConstruct
    public void start()
    {
        try {
            cassandra.start();
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }

        waitForListener();
    }

    void waitForListener()
    {
        InetSocketAddress testAddress = new InetSocketAddress(rpcAddress, rpcPort);
        int remainingMilliseconds = 5000;
        final long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(remainingMilliseconds);
        do {
            try {
                Socket testSocket = new Socket();
                testSocket.connect(testAddress, remainingMilliseconds);
                testSocket.close();
                return;
            }
            catch (IOException e) {
                try {
                    TimeUnit.MILLISECONDS.sleep(100);
                }
                catch (InterruptedException interrupted) {
                    Thread.currentThread().interrupt();
                    throw new RuntimeException(interrupted);
                }
            }
            remainingMilliseconds = (int)(TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime()));
        } while (remainingMilliseconds > 0);
    }

    @PreDestroy
    public void stop()
    {
        cassandra.stop();
    }

    public InetAddress getRpcAddress()
    {
        return rpcAddress;
    }

    public int getRpcPort()
    {
        return rpcPort;
    }
}

@@ -330,16 +330,54 @@
 deInfo;%0A
+import com.proofpoint.units.Duration;%0A
 import o
@@ -4678,24 +4678,25 @@
 cPort);%0A
+%0A
         
 int rema
@@ -4691,91 +4691,95 @@
     
-int remainingMilliseconds = 5000;%0A        final long deadline = System.nanoTime() +
+long start = System.nanoTime();%0A%0A        while (Duration.nanosSince(start).toMillis() %3C
  Tim
@@ -4790,61 +4790,26 @@
 t.MI
-LLISECONDS.toNanos(remainingMilliseconds);%0A        do
+NUTES.toMillis(1))
  %7B%0A 
@@ -4923,37 +4923,50 @@
 ddress, 
-remainingMilliseconds
+(int) TimeUnit.SECONDS.toMillis(1)
 );%0A     
@@ -4992,24 +4992,25 @@
 et.close();%0A
+%0A
             
@@ -5126,13 +5126,8 @@
 nit.
-MILLI
 SECO
@@ -5137,18 +5137,16 @@
 .sleep(1
-00
 );%0A     
@@ -5378,148 +5378,9 @@
     
-    remainingMilliseconds = (int)(TimeUnit.NANOSECONDS.toMillis(deadline - System.nanoTime()));%0A        %7D while (remainingMilliseconds %3E 0);
+%7D
 %0A   

Make the code easier to read