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

import org.apache.tinkerpop.gremlin.structure.*;

import com.orientechnologies.orient.core.db.record.OIdentifiable;
import com.orientechnologies.orient.core.record.impl.ODocument;

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

import org.apache.commons.lang.NotImplementedException;


public class OrientElement implements Element {
    protected OIdentifiable rawElement;
    protected OrientGraph graph;

    public OrientElement(final OrientGraph graph, final OIdentifiable rawElement) {
        this.graph = graph;
        this.rawElement = rawElement;
    }

    public Object id() {
        return rawElement.getIdentity();
    }

    public String label() {
        throw new NotImplementedException();
    }

    public Graph graph() {
        return graph;
    }

    public <V> Property<V> property(final String key, final V value) {
        if (rawElement instanceof ODocument) {
            ODocument doc = (ODocument) rawElement;
            doc.field(key, value);
            doc.save();
        }

        return new OrientProperty<>(key, value, this);
    }

    public void remove() {
        throw new NotImplementedException();
    }

    public <V> Iterator<? extends Property<V>> properties(final String... propertyKeys) {
        ODocument raw = rawElement.getRecord();
        Map<String, Object> properties = raw.toMap();
        HashSet<String> keys = new HashSet<>(Arrays.asList(propertyKeys));

        Stream<Map.Entry<String, Object>> entries = StreamUtils.asStream(properties.entrySet().iterator());
        if (keys.size() > 0) entries = entries.filter(entry -> keys.contains(entry));

        Stream<OrientProperty<V>> propertyStream = entries.map(entry -> new OrientProperty<>(entry.getKey(), (V) entry.getValue(), this));
        return propertyStream.iterator();
    }

}

@@ -51,16 +51,128 @@
 cture;%0A%0A
+import com.orientechnologies.orient.core.id.ORecordId;%0Aimport com.orientechnologies.orient.core.record.ORecord;%0A
 import o
@@ -1025,16 +1025,18 @@
      if 
+(!
 (rawElem
@@ -1056,31 +1056,93 @@
 f ODocument)
- %7B
+)
 %0A    
+        rawElement = new ODocument(rawElement.getIdentity());%0A%0A
         ODoc
@@ -1177,20 +1177,16 @@
 lement;%0A
-    
         
@@ -1212,28 +1212,24 @@
 e);%0A        
-    
 doc.save();%0A
@@ -1232,19 +1232,8 @@
 ();%0A
-        %7D%0A%0A
     
@@ -1478,18 +1478,21 @@
 cument r
-aw
+ecord
  = rawEl
@@ -1506,24 +1506,91 @@
 etRecord();%0A
+        if (record == null)%0A            record = new ODocument();%0A%0A
         Map%3C
@@ -1619,18 +1619,21 @@
 ties = r
-aw
+ecord
 .toMap()

support elements if not loaded yet, too - jdoesn't work yet ;(