/* Copyright 2014 The Johns Hopkins University Applied Physics Laboratory
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.jhuapl.tinkerpop;

import java.util.Iterator;
import java.util.Map.Entry;

import org.apache.accumulo.core.client.BatchWriter;
import org.apache.accumulo.core.client.MutationsRejectedException;
import org.apache.accumulo.core.client.Scanner;
import org.apache.accumulo.core.client.ScannerBase;
import org.apache.accumulo.core.data.Key;
import org.apache.accumulo.core.data.Mutation;
import org.apache.accumulo.core.data.Range;
import org.apache.accumulo.core.data.Value;
import org.apache.accumulo.core.util.PeekingIterator;
import org.apache.hadoop.io.Text;

import com.tinkerpop.blueprints.CloseableIterable;
import com.tinkerpop.blueprints.Edge;
import com.tinkerpop.blueprints.Element;
import com.tinkerpop.blueprints.Index;

import edu.jhuapl.tinkerpop.tables.NamedIndexTableWrapper;

/**
 * Accumulo-based index implementation
 * @param <T>
 */
public class AccumuloIndex<T extends Element> implements Index<T> {
  private final GlobalInstances globals;
  private final Class<T> indexedType;
  private final String indexName;
  private final NamedIndexTableWrapper indexWrapper;

  public AccumuloIndex(GlobalInstances globals, String indexName, Class<T> indexedType) {
    this.globals = globals;
    this.indexName = indexName;
    this.indexedType = indexedType;

    try {
      if (!globals.getConfig().getConnector().tableOperations().exists(getTableName())) {
        globals.getConfig().getConnector().tableOperations().create(getTableName());
      }
    } catch (Exception e) {
      throw new AccumuloGraphException(e);
    }

    this.indexWrapper = new NamedIndexTableWrapper(globals,
        indexedType, indexName);
  }

  @Override
  public String getIndexName() {
    return indexName;
  }

  public String getTableName() {
    return globals.getConfig().getNamedIndexTableName(indexName);
  }

  @Override
  public Class<T> getIndexClass() {
    return indexedType;
  }

  @Override
  public void put(String key, Object value, Element element) {
    element.setProperty(key, value);
    Mutation m = new Mutation(AccumuloByteSerializer.serialize(value));
    m.put(key.getBytes(), element.getId().toString().getBytes(), "".getBytes());
    BatchWriter w = getWriter();
    try {
      w.addMutation(m);
      w.flush();
    } catch (MutationsRejectedException e) {
      throw new AccumuloGraphException(e);
    }
  }

  @Override
  public CloseableIterable<T> get(String key, Object value) {
    Scanner scan = getScanner();
    byte[] id = AccumuloByteSerializer.serialize(value);
    scan.setRange(Range.exact(new Text(id)));
    scan.fetchColumnFamily(new Text(key));
    return new IndexIterable(globals.getGraph(), scan, indexedType);
  }

  @Override
  public CloseableIterable<T> query(String key, Object query) {
    throw new UnsupportedOperationException();
  }

  @Override
  public long count(String key, Object value) {
    CloseableIterable<T> iterable = get(key, value);
    Iterator<T> iter = iterable.iterator();
    int count = 0;
    while (iter.hasNext()) {
      count++;
      iter.next();
    }
    iterable.close();
    return count;
  }

  @Override
  public void remove(String key, Object value, Element element) {
    Mutation m = new Mutation(AccumuloByteSerializer.serialize(value));
    m.putDelete(key.getBytes(), element.getId().toString().getBytes());
    BatchWriter w = getWriter();
    try {
      w.addMutation(m);
      w.flush();
    } catch (MutationsRejectedException e) {
      throw new AccumuloGraphException(e);
    }

  }

  private BatchWriter getWriter() {
    return globals.getGraph().getWriter(getTableName());
  }

  private Scanner getScanner() {
    return globals.getGraph().getScanner(getTableName());
  }

  private class IndexIterable implements CloseableIterable<T> {
    private ScannerBase scan;
    private Class<T> indexedType;

    private IndexIterable(AccumuloGraph parent, ScannerBase scan,
        Class<T> indexedType) {
      this.scan = scan;
      this.indexedType = indexedType;
    }

    @Override
    public Iterator<T> iterator() {
      if (scan != null) {
        return new ScannerIterable<T>(scan) {

          @SuppressWarnings("unchecked")
          @Override
          public T next(PeekingIterator<Entry<Key, Value>> iterator) {
            String id = iterator.next()
                .getKey().getColumnQualifier().toString();
            // TODO better use of information readily
            // available...
            if (indexedType.equals(Edge.class)) {
              return (T) new AccumuloEdge(globals, id);
            }
            else {
              return (T) new AccumuloVertex(globals, id);
            }
          }
        }.iterator();
      }
      else {
        return null;
      }
    }

    @Override
    public void close() {
      if (scan != null) {
        scan.close();
        scan = null;
      }
    }
  }
}

@@ -660,16 +660,41 @@
 erpop;%0A%0A
+import java.util.Arrays;%0A
 import j
@@ -1395,16 +1395,217 @@
 .Index;%0A
+import com.tinkerpop.blueprints.Vertex;%0A%0Aimport edu.jhuapl.tinkerpop.parser.EdgeIndexParser;%0Aimport edu.jhuapl.tinkerpop.parser.ElementIndexParser;%0Aimport edu.jhuapl.tinkerpop.parser.VertexIndexParser;
 %0Aimport 
@@ -4897,24 +4897,212 @@
 terator() %7B%0A
+      final ElementIndexParser%3C? extends AccumuloElement%3E parser =%0A          Vertex.class.equals(indexedType) ? new VertexIndexParser(globals) :%0A            new EdgeIndexParser(globals);%0A%0A
       if (sc
@@ -5310,387 +5310,64 @@
     
-String id = iterator.next()%0A                .getKey().getColumnQualifier().toString();%0A            // TODO better use of information readily%0A            // available...%0A            if (indexedType.equals(Edge.class)) %7B%0A              return (T) new AccumuloEdge(globals, id);%0A            %7D%0A            else %7B%0A              return (T) new AccumuloVertex(globals, id);%0A            %7D
+return (T) parser.parse(Arrays.asList(iterator.next()));
 %0A   

Use index parser in AccumuloIndex