package org.apache.tinkerpop.gremlin.elastic.process.graph.traversal.steps;

import org.apache.tinkerpop.gremlin.elastic.elasticservice.ElasticService;
import org.apache.tinkerpop.gremlin.process.traversal.step.sideEffect.GraphStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.util.HasContainer;
import org.apache.tinkerpop.gremlin.structure.*;

import java.util.*;

public class ElasticGraphStep<E extends Element> extends GraphStep<E> {

    private final ArrayList<HasContainer> hasContainers;
    private final ElasticService elasticService;
    private Integer resultLimit;

    public ElasticGraphStep(GraphStep originalStep, ArrayList<HasContainer> hasContainers, ElasticService elasticService,Integer resultLimit) {
        super(originalStep.getTraversal(),originalStep.getReturnClass(),originalStep.getIds());
        if (originalStep.getLabel().isPresent()) this.setLabel(originalStep.getLabel().get().toString());
        this.hasContainers = hasContainers;
        this.elasticService = elasticService;
        this.resultLimit = resultLimit;
        this.setIteratorSupplier(() -> (Iterator<E>) (Vertex.class.isAssignableFrom(this.returnClass) ? this.vertices() : this.edges()));
    }

    private Iterator<? extends Vertex> vertices() {
        ArrayList<HasContainer> hasList = hasContainers;
        Object[] ids = super.getIds();
        if(ids.length > 0) {
            hasList  = (ArrayList<HasContainer>) hasContainers.clone();
            hasList.add(new HasContainer("~id", Contains.within, ids));
        }
        return elasticService.searchVertices(hasList, resultLimit);
    }

    private Iterator<? extends Edge> edges() {
        ArrayList<HasContainer> hasList = hasContainers;
        Object[] ids = super.getIds();
        if(ids.length > 0) {
            hasList  = (ArrayList<HasContainer>) hasContainers.clone();
            hasList.add(new HasContainer("~id", Contains.within, ids));
        }
         return elasticService.searchEdges(hasList, resultLimit, null, null);
    }
}

@@ -145,16 +145,155 @@
 ervice;%0A
+import org.apache.tinkerpop.gremlin.process.traversal.Traverser;%0Aimport org.apache.tinkerpop.gremlin.process.traversal.TraverserGenerator;%0A
 import o
@@ -439,24 +439,125 @@
 sContainer;%0A
+import org.apache.tinkerpop.gremlin.process.traversal.traverser.B_O_P_PA_S_SE_SL_TraverserGenerator;%0A
 import org.a
@@ -2265,10 +2265,1323 @@
 ;%0A    %7D%0A
+%0A%0A    public void generateTraversers(final TraverserGenerator traverserGenerator) %7B%0A%0A        try %7B%0A            this.start = this.iteratorSupplier.get();%0A            if (null != this.start) %7B%0A%0A                if (this.start instanceof Iterator) %7B%0A                    List%3CE%3E newListForIterator = new ArrayList%3C%3E();%0A                    Iterator%3CE%3E iter = (Iterator%3CE%3E) this.start;%0A                    while(iter.hasNext())%7B%0A                        E next = iter.next();%0A%0A                        this.starts.add(B_O_P_PA_S_SE_SL_TraverserGenerator.instance().generate(next,this,1l));%0A                        newListForIterator.add(next);%0A                    %7D%0A                    this.start = newListForIterator.iterator();%0A%0A                %7D else %7B%0A                    this.starts.add(traverserGenerator.generate((E) this.start, this, 1l));%0A                %7D%0A            %7D%0A        %7Dcatch (NoSuchElementException ex)%7B%0A            throw ex;%0A        %7D%0A        catch (final Exception e) %7B%0A            throw new IllegalStateException(e.getMessage(), e);%0A        %7D%0A    %7D%0A%0A    @Override%0A    protected Traverser%3CE%3E processNextStart() %7B%0A        if (this.first) %7B%0A            generateTraversers(this.getTraversal().getTraverserGenerator());%0A            this.first = false;%0A        %7D%0A        return this.starts.next();%0A    %7D%0A%0A
 %7D%0A

fix Where and select tests