package io.atomix.protocols.raft.test;

import io.atomix.cluster.MemberId;
import io.atomix.cluster.Node;
import io.atomix.cluster.discovery.BootstrapDiscoveryProvider;
import io.atomix.cluster.discovery.NodeDiscoveryProvider;
import io.atomix.core.Atomix;
import io.atomix.core.list.AsyncDistributedList;
import io.atomix.protocols.raft.partition.RaftPartitionGroup;
import io.atomix.storage.StorageLevel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.util.concurrent.CountDownLatch;

public class LatencyTest {
  public static void main(String[] args) throws InterruptedException {
    new LatencyTest().run();
  }

  private final MemberId member = MemberId.anonymous();

  private void run() throws InterruptedException {
    final Atomix atomix = startAtomix();
    atomix.start().join();

    final Logger logger = LoggerFactory.getLogger(this.getClass());
    final AsyncDistributedList<Integer> list = atomix.<Integer>getList("list").async();
    final int workCount = 15_000;
    final CountDownLatch latch = new CountDownLatch(workCount);
    for (int i = 0; i < workCount; i++) {
      list.add(i, i)
          .whenComplete(
              (nothing, error) -> {
                if (latch.getCount() % 100 == 0) {
                  logger.info("Count down: {}", latch.getCount());
                }

                latch.countDown();
              });
    }

    latch.await();
    atomix.stop().join();
  }

  private Atomix startAtomix() {
    final RaftPartitionGroup system =
        RaftPartitionGroup.builder("system")
            .withMembers(member)
            .withNumPartitions(1)
            .withPartitionSize(1)
            .withDataDirectory(new File(String.format("target/latency-logs/%s/system", member.id())))
            .build();
    final RaftPartitionGroup data =
        RaftPartitionGroup.builder("data")
            .withMembers(member)
            .withNumPartitions(1)
            .withPartitionSize(1)
            .withStorageLevel(StorageLevel.DISK)
            .withFlushOnCommit(false)
            .withDataDirectory(new File(String.format("target/latency-logs/%s/data", member.id())))
            .build();
    final NodeDiscoveryProvider discoveryProvider =
        BootstrapDiscoveryProvider.builder()
            .withNodes(Node.builder().withId(member.id()).build())
            .build();

    return Atomix.builder()
        .withMemberId(member)
        .withShutdownHookEnabled()
        .withMembershipProvider(discoveryProvider)
        .withManagementGroup(system)
        .withPartitionGroups(data)
        .build();
  }
}