package com.tinkerpop.gremlin.structure.strategy;

import com.tinkerpop.gremlin.structure.Edge;
import com.tinkerpop.gremlin.structure.Element;
import com.tinkerpop.gremlin.structure.Graph;
import com.tinkerpop.gremlin.structure.Property;
import com.tinkerpop.gremlin.structure.Strategy;
import com.tinkerpop.gremlin.structure.Vertex;
import com.tinkerpop.gremlin.structure.util.ElementHelper;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.function.UnaryOperator;

/**
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */
public class IdGraphStrategy implements GraphStrategy {

    private final String idKey;

    private final Supplier<?> edgeIdSupplier;
    private final Supplier<?> vertexIdSupplier;

    private final boolean supportsVertexId;
    private final boolean supportsEdgeId;

    private IdGraphStrategy(final String idKey, final Supplier<?> vertexIdSupplier,
                            final Supplier<?> edgeIdSupplier, final boolean supportsVertexId,
                            final boolean supportsEdgeId) {
        // assumes this is an indexed key.
        this.idKey = Property.Key.hidden(idKey);

        this.edgeIdSupplier = edgeIdSupplier;
        this.vertexIdSupplier = vertexIdSupplier;
        this.supportsEdgeId = supportsEdgeId;
        this.supportsVertexId = supportsVertexId;
    }

    @Override
    public UnaryOperator<Function<Object[], Vertex>> getAddVertexStrategy(final Strategy.Context<Graph> ctx) {
        return (f) -> (keyValues) -> {
            final List<Object> o = new ArrayList<>(Arrays.asList(keyValues));
            if (supportsVertexId) {
                final Object val = ElementHelper.getIdValue(keyValues).orElse(vertexIdSupplier.get());
                final int pos = o.indexOf(Element.ID);
                if (pos > -1) {
                    o.remove(pos);
                    o.remove(pos);
                }

                o.addAll(Arrays.asList(this.idKey, val));
            }

            return f.apply(o.toArray());
        };
    }

    @Override
    public UnaryOperator<Function<Object, Vertex>> getGraphvStrategy(final Strategy.Context<Graph> ctx) {
        // don't apply f because the implementation needs to be highjacked by the Strategy
        if (supportsVertexId)
            return (f) -> (id) -> (Vertex) ctx.getGraph().V().has(idKey, id).next();
        else
            return UnaryOperator.identity();
    }

    @Override
    public UnaryOperator<Function<Object, Edge>> getGrapheStrategy(final Strategy.Context<Graph> ctx) {
        // don't apply f because the implementation needs to be highjacked by the Strategy
        return (f) -> (id) -> (Edge) ctx.getGraph().E().has(idKey, id).next();
    }

    @Override
    public UnaryOperator<Supplier<Object>> getElementGetId(final Strategy.Context<? extends Element> ctx) {
        // if the property is not present then it's likely an internal call from the graph on addVertex in which case,
        // the base implementation should be called
        return (f) -> () -> ctx.getCurrent().getProperty(idKey).orElse(f.get());
    }

    public static final class Builder {
        private final String idKey;
        private Supplier<?> vertexIdSupplier;
        private Supplier<?> edgeIdSupplier;
        private boolean supportsVertexId;
        private boolean supportsEdgeId;

        public Builder(final String idKey) {
            this.idKey = idKey;
            this.edgeIdSupplier = this::supplyStringId;
            this.vertexIdSupplier = this::supplyStringId;
            this.supportsEdgeId = true;
            this.supportsVertexId = true;
        }

        public IdGraphStrategy build() {
            if (!this.supportsEdgeId && !this.supportsVertexId)
                throw new IllegalStateException("Since supportsEdgeId and supportsVertexId are false, there is no need to use IdGraphStrategy");

            return new IdGraphStrategy(this.idKey, this.vertexIdSupplier, this.edgeIdSupplier,
                    this.supportsVertexId, this.supportsEdgeId);
        }

        public Builder vertexIdMaker(final Supplier<?> vertexIdSupplier) {
            if (null == vertexIdSupplier)
                throw new IllegalArgumentException("vertexIdSupplier");

            this.vertexIdSupplier = vertexIdSupplier;
            return this;
        }

        public Builder edgeIdMaker(final Supplier<?> edgeIdSupplier) {
            if (null == edgeIdSupplier)
                throw new IllegalArgumentException("edgeIdSupplier");

            this.edgeIdSupplier = edgeIdSupplier;
            return this;
        }

        public Builder supportsEdgeId(final boolean supports) {
            this.supportsEdgeId = supports;
            return this;
        }

        public Builder supportsVertexId(final boolean supports) {
            this.supportsVertexId = supports;
            return this;
        }

        private String supplyStringId() {
            return UUID.randomUUID().toString();
        }
    }
}

@@ -386,16 +386,72 @@
 tHelper;
+%0Aimport com.tinkerpop.gremlin.util.function.TriFunction;
 %0A%0Aimport
@@ -1692,526 +1692,369 @@
  -%3E 
-%7B%0A            final List%3CObject%3E o = new ArrayList%3C%3E(Arrays.asList(keyValues));%0A            if (supportsVertexId) %7B%0A                final Object val = ElementHelper.getIdValue(keyValues).orElse(vertexIdSupplier.get());%0A                final int pos = o.indexOf(Element.ID);%0A                if (pos %3E -1) %7B%0A                    o.remove(pos);%0A                    o.remove(pos);%0A                %7D%0A%0A                o.addAll(Arrays.asList(this.idKey, val));%0A            %7D%0A%0A            return f.apply(o.toArray());%0A        %7D
+f.apply(this.injectId(supportsVertexId, keyValues, vertexIdSupplier).toArray());%0A    %7D%0A%0A    @Override%0A    public UnaryOperator%3CTriFunction%3CString, Vertex, Object%5B%5D, Edge%3E%3E getAddEdgeStrategy(final Strategy.Context%3CVertex%3E ctx) %7B%0A        return (f) -%3E (label, v, keyValues) -%3E f.apply(label, v, this.injectId(supportsEdgeId, keyValues, edgeIdSupplier).toArray())
 ;%0A  
@@ -3121,24 +3121,571 @@
 ());%0A    %7D%0A%0A
+    private List%3CObject%3E injectId(final boolean supports, final Object%5B%5D keyValues, final Supplier%3C?%3E idMaker) %7B%0A        final List%3CObject%3E o = new ArrayList%3C%3E(Arrays.asList(keyValues));%0A        if (supports) %7B%0A            final Object val = ElementHelper.getIdValue(keyValues).orElse(idMaker.get());%0A            final int pos = o.indexOf(Element.ID);%0A            if (pos %3E -1) %7B%0A                o.remove(pos);%0A                o.remove(pos);%0A            %7D%0A%0A            o.addAll(Arrays.asList(this.idKey, val));%0A        %7D%0A%0A        return o;%0A    %7D%0A%0A
     public s

Refactor IdGraphStrategy to extract reusable function.