package org.apache.tinkerpop.gremlin.orientdb.structure;

import com.orientechnologies.orient.core.id.ORecordId;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.lang.NotImplementedException;
import org.apache.tinkerpop.gremlin.process.computer.GraphComputer;
import org.apache.tinkerpop.gremlin.process.traversal.TraversalSource;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.apache.tinkerpop.gremlin.structure.Edge;
import org.apache.tinkerpop.gremlin.structure.Graph;
import org.apache.tinkerpop.gremlin.structure.Transaction;
import org.apache.tinkerpop.gremlin.structure.Vertex;
import org.apache.tinkerpop.gremlin.structure.io.Io;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Stream;

import com.orientechnologies.orient.core.record.ORecord;
import com.orientechnologies.orient.core.iterator.ORecordIteratorClass;
import com.orientechnologies.orient.core.db.document.ODatabaseDocumentTx;
import com.orientechnologies.orient.core.db.record.OIdentifiable;
import com.orientechnologies.orient.core.record.impl.ODocument;
import com.orientechnologies.orient.core.record.impl.ODocumentInternal;

import static org.apache.tinkerpop.gremlin.orientdb.structure.StreamUtils.asStream;


public final class OrientGraph implements Graph {

    public static final Logger LOGGER = LoggerFactory.getLogger(OrientGraph.class);
    protected ODatabaseDocumentTx database;

    public OrientGraph(ODatabaseDocumentTx iDatabase) {
        this.database = iDatabase;
    }

    @Override
    public Vertex addVertex(Object... keyValues) {
        throw new NotImplementedException();
    }

    @Override
    public <C extends GraphComputer> C compute(Class<C> graphComputerClass) throws IllegalArgumentException {
        throw new NotImplementedException();
    }

    @Override
    public GraphComputer compute() throws IllegalArgumentException {
        throw new NotImplementedException();
    }

    @Override
    public Iterator<Vertex> vertices(Object... vertexIds) {
        //TODO: limit on vertexIds
        boolean polymorphic = true;
        // String iClassName = OrientVertexType.CLASS_NAME;
        String elementClass = "V";
        OrientGraph graph = this;

        Iterator<ORecord> itty = new ORecordIteratorClass<ORecord>(database, database, elementClass, polymorphic);
        Stream<Vertex> stream = asStream(itty).map(r -> toVertex(r));

        if (vertexIds.length > 0) {
            //TODO: use an index for the lookup instead of scanning all vertices!
            HashSet<Object> vertexIdsSet = new HashSet<Object>(Arrays.asList(vertexIds));
            stream = stream.filter(v -> vertexIdsSet.contains(v.id()) || vertexIdsSet.contains(v.id().toString()));
        }

        return stream.iterator();
    }

    private Vertex toVertex(ORecord record) {
        // taken mostly from old blueprints driver, could be done nicer I believe
        OrientElement currentElement = null;

        if (record == null) throw new NoSuchElementException();

        if (record instanceof OIdentifiable)
            record = record.getRecord();

        if (record instanceof ODocument) {
            final ODocument currentDocument = (ODocument) record;

            if (currentDocument.getInternalStatus() == ODocument.STATUS.NOT_LOADED)
                currentDocument.load();

            if (ODocumentInternal.getImmutableSchemaClass(currentDocument) == null)
                throw new IllegalArgumentException(
                    "Cannot determine the graph element type because the document class is null. Probably this is a projection, use the EXPAND() function");

            // if (ODocumentInternal.getImmutableSchemaClass(currentDocument).isSubClassOf(graph.getEdgeBaseType()))
            //   currentElement = new OrientEdge(graph, currentDocument);
            // else
            currentElement = new OrientVertex(this, currentDocument);
        }

        return (Vertex) currentElement;
    }

    @Override
    public Iterator<Edge> edges(Object... edgeIds) {
        throw new NotImplementedException();
    }

    @Override
    public Transaction tx() {
        throw new NotImplementedException();
    }

    @Override
    public Variables variables() {
        throw new NotImplementedException();
    }

    @Override
    public Configuration configuration() {
        throw new NotImplementedException();
    }

    @Override
    public void close() throws Exception {
    }
}

@@ -2298,24 +2298,25 @@
  = %22V%22;%0A
+%0A
         
 OrientGr
@@ -2311,35 +2311,86 @@
     
-OrientGraph graph = this;%0A%0A
+if (vertexIds.length == 0) %7B%0A            // return all vertices as stream%0A    
     
@@ -2443,23 +2443,16 @@
 orClass%3C
-ORecord
 %3E(databa
@@ -2505,31 +2505,18 @@
     
-Stream%3CVertex%3E stream =
+    return
  asS
@@ -2552,320 +2552,119 @@
 (r))
-;%0A%0A        if (vertexIds.length %3E 0) %7B%0A            //TODO: use an index for the lookup instead of scanning all vertices!%0A            HashSet%3CObject%3E vertexIdsSet = new HashSet%3CObject%3E(Arrays.asList(vertexIds));%0A            stream = stream.filter(v -%3E vertexIdsSet.contains(v.id()) %7C%7C vertexIdsSet.contains(v.id()
+.iterator();%0A        %7D else %7B%0A            Stream%3CORecordId%3E ids = Stream.of(vertexIds).map(v -%3E new ORecordId(v
 .toS
@@ -2674,31 +2674,24 @@
 ng()));%0A    
-    %7D%0A%0A
         retu
@@ -2693,22 +2693,66 @@
  return 
-stream
+ids.map(id -%3E (Vertex) new OrientVertex(this, id))
 .iterato
@@ -2752,24 +2752,34 @@
 iterator();%0A
+        %7D%0A
     %7D%0A%0A    p
@@ -4445,20 +4445,65 @@
 Exception %7B%0A
+        throw new NotImplementedException();%0A
     %7D%0A%7D%0A

use index to lookup vertices by id