package com.netflix.ndbench.plugins.janusgraph;

import com.netflix.ndbench.api.plugin.DataGenerator;
import com.netflix.ndbench.api.plugin.NdBenchClient;
import com.netflix.ndbench.api.plugin.annotations.NdBenchClientPlugin;
import com.netflix.ndbench.plugins.janusgraph.configs.IJanusGraphConfig;
import com.netflix.ndbench.plugins.janusgraph.cql.JanusGraphBuilderCQLProvider;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.apache.tinkerpop.gremlin.structure.T;
import org.apache.tinkerpop.gremlin.structure.Transaction;
import org.apache.tinkerpop.gremlin.structure.Vertex;
import org.janusgraph.core.JanusGraph;
import org.janusgraph.core.JanusGraphFactory;
import org.janusgraph.core.JanusGraphTransaction;
import org.janusgraph.core.JanusGraphVertex;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.inject.Inject;
import javax.inject.Singleton;
import java.util.Iterator;
import java.util.List;

/***
 * JanusGraph benchmarking plugin to measure throughput of write and read by
 * single ID using JanusGraph Core API
 *
 * @author pencal
 */
@Singleton
@NdBenchClientPlugin("janusgraph-cql")
public class JanusGraphPluginCQL extends JanusGraphBasePlugin implements NdBenchClient {
    private static final Logger Logger = LoggerFactory.getLogger(JanusGraphPluginCQL.class);
    private static String BACKEND = "cql";

    private final JanusGraphFactory.Builder graphBuilder;

    private DataGenerator dataGenerator;
    private GraphTraversalSource traversalSource;
    private JanusGraph graph;
    private boolean useJanusgraphTransaction;

    @Inject
    public JanusGraphPluginCQL(IJanusGraphConfig config, JanusGraphBuilderCQLProvider builderProvider) {
        super(BACKEND, config.getStorageHostname(), config.getStoragePort());
        this.graphBuilder = builderProvider.getGraphBuilder();
        this.useJanusgraphTransaction = config.useJanusgraphTransaction();
    }

    @Override
    public void init(DataGenerator dataGenerator) throws Exception {
        this.graph = graphBuilder.open();
        this.traversalSource = graph.traversal();
        this.dataGenerator = dataGenerator;
        createSchema(graph);
    }

    @Override
    public String readSingle(String key) throws Exception {
        String response = OK;
        if (useJanusgraphTransaction) {
            try (JanusGraphTransaction tx = graph.newTransaction()) {
                Iterator<JanusGraphVertex> result = tx.query().has(PROP_CUSTOM_ID_KEY, key).vertices().iterator();
                if (result == null)
                    throw new Exception(
                            "Internal error when reading data with key" + key + " using JanusGraph Core API");
                else if (!result.hasNext())
                    response = CACHE_MISS;
            }
        } else {
            List<Vertex> results = traversalSource.V().has(PROP_CUSTOM_ID_KEY, key).toList();

            if (results == null)
                throw new Exception("Internal error when reading data with key" + key + " using TinkerPop API");
            else if (results.size() == 0)
                response = CACHE_MISS;
        }

        return response;
    }

    @Override
    public String writeSingle(String key) throws Exception {
        if (useJanusgraphTransaction) {
            try (JanusGraphTransaction tx = graph.newTransaction()) {
                tx.addVertex(T.label, VERTEX_LABEL_LEVEL_1, PROP_CUSTOM_ID_KEY, key, PROP_METADATA_KEY,
                        dataGenerator.getRandomValue());
                tx.commit();
            }
        } else {
            Transaction tx = traversalSource.tx();
            tx.open();
            traversalSource.getGraph().addVertex(T.label, VERTEX_LABEL_LEVEL_1, PROP_CUSTOM_ID_KEY, key,
                    PROP_METADATA_KEY, dataGenerator.getRandomValue());
            tx.commit();
        }

        return OK;
    }

    @Override
    public void shutdown() throws Exception {
        graph.close();
    }

    @Override
    public String getConnectionInfo() throws Exception {
        return super.getConnectionInfo(graph);
    }

    @Override
    public String runWorkFlow() throws Exception {
        return null;
    }
}

@@ -2446,17 +2446,8 @@
     
-Iterator%3C
 Janu
@@ -2462,18 +2462,36 @@
 rtex
-%3E result =
+ vertex = (JanusGraphVertex)
  tx.
@@ -2541,20 +2541,23 @@
 es()
-.iterator();
+;              
 %0A   
@@ -2573,22 +2573,22 @@
     if (
-result
+vertex
  == null
@@ -2580,32 +2580,67 @@
 (vertex == null)
+ %7B%0A                    tx.commit();
 %0A               
@@ -2796,32 +2796,50 @@
     
-else if (!result.hasNext
+%7D%0A                if(vertex.keys().isEmpty
 ())%0A
@@ -2885,33 +2885,81 @@
 SS;%0A            
-%7D
+    %0A                tx.commit();%0A            %7D %0A
 %0A        %7D else 
@@ -3456,84 +3456,17 @@
     
-try (JanusGraphTransaction tx = graph.newTransaction()) %7B%0A                tx
+    graph
 .add
@@ -3595,32 +3595,72 @@
 tRandomValue());
+ //Automatically opens a new transaction
 %0A               
@@ -3652,34 +3652,42 @@
                 
-tx
+graph.tx()
 .commit();%0A     
@@ -3693,104 +3693,16 @@
     
-    %7D%0A        %7D else %7B%0A            Transaction tx = traversalSource.tx();%0A            tx.open();
+%7D else %7B
 %0A   
@@ -3880,33 +3880,62 @@
 );%0A            t
-x
+raversalSource.getGraph().tx()
 .commit();%0A     
@@ -3957,33 +3957,32 @@
 eturn OK;%0A    %7D%0A
-%0A
     @Override%0A  

Simplify Janusgrph transactions and APIs