package io.lumify.palantir.dataImport;

import com.google.common.base.Function;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.Maps;
import io.lumify.core.util.LumifyLogger;
import io.lumify.core.util.LumifyLoggerFactory;
import org.securegraph.Vertex;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public abstract class PtImporterBase<T> {
    private static final LumifyLogger LOGGER = LumifyLoggerFactory.getLogger(PtImporterBase.class);
    private final Class<T> ptClass;
    private final DataImporter dataImporter;
    private final LoadingCache<String, Vertex> vertexCache;

    protected PtImporterBase(final DataImporter dataImporter, Class<T> ptClass) {
        this.dataImporter = dataImporter;
        this.ptClass = ptClass;

        vertexCache = CacheBuilder.newBuilder()
                .maximumSize(1000)
                .build(new CacheLoader<String, Vertex>() {
                    @Override
                    public Vertex load(String key) throws Exception {
                        return dataImporter.getGraph().getVertex(key, dataImporter.getAuthorizations());
                    }
                });
    }

    public void run() {
        LOGGER.info(this.getClass().getName());
        Iterable<T> rows = dataImporter.getSqlRunner().select(getSql(), ptClass);
        beforeProcessRows();
        long count = run(rows);
        afterProcessRows();
        LOGGER.info("Imported %d %s", count, this.ptClass.getSimpleName());
    }

    protected abstract long run(Iterable<T> rows);

    protected void afterProcessRows() {

    }

    protected void beforeProcessRows() {

    }

    protected DataImporter getDataImporter() {
        return dataImporter;
    }

    protected abstract String getSql();

    public LoadingCache<String, Vertex> getVertexCache() {
        return vertexCache;
    }

    protected Map<String, Vertex> getVertexCacheVertices(Iterable<String> vertexIds) {
        List<String> vertexIdsToGet = new ArrayList<String>();
        List<Vertex> vertices = new ArrayList<Vertex>();

        // get the vertices already in the cache
        for (String vertexId : vertexIds) {
            Vertex v = this.vertexCache.getIfPresent(vertexId);
            if (v == null) {
                vertexIdsToGet.add(vertexId);
            } else {
                vertices.add(v);
            }
        }

        // get the vertices not in the cache
        Iterable<Vertex> newVertices = getDataImporter().getGraph().getVertices(vertexIdsToGet, getDataImporter().getAuthorizations());
        for (Vertex v : newVertices) {
            vertices.add(v);
            this.vertexCache.put(v.getId(), v);
        }

        // build a map to return
        return Maps.uniqueIndex(vertices, new Function<Vertex, String>() {
            @Override
            public String apply(Vertex input) {
                return input.getId();
            }
        });
    }

    protected String getObjectVertexId(long objectId) {
        return getDataImporter().getIdPrefix() + objectId;
    }

    public Class<T> getPtClass() {
        return ptClass;
    }
}

@@ -1298,24 +1298,77 @@
 oid run() %7B%0A
+        long startTime = System.currentTimeMillis();%0A
         LOGG
@@ -1566,32 +1566,83 @@
 rProcessRows();%0A
+        long endTime = System.currentTimeMillis();%0A
         LOGGER.i
@@ -1660,16 +1660,29 @@
 ed %25d %25s
+ (time: %25dms)
 %22, count
@@ -1707,24 +1707,45 @@
 SimpleName()
+, endTime - startTime
 );%0A    %7D%0A%0A  

add time to results