/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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 org.apache.atlas.repository.graphdb.janus;

import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasErrorCode;
import org.apache.atlas.AtlasException;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.groovy.GroovyExpression;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.repository.graphdb.AtlasEdge;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasGraphIndexClient;
import org.apache.atlas.repository.graphdb.AtlasGraphManagement;
import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
import org.apache.atlas.repository.graphdb.AtlasGraphTraversal;
import org.apache.atlas.repository.graphdb.AtlasIndexQuery;
import org.apache.atlas.repository.graphdb.AtlasIndexQueryParameter;
import org.apache.atlas.repository.graphdb.AtlasPropertyKey;
import org.apache.atlas.repository.graphdb.AtlasSchemaViolationException;
import org.apache.atlas.repository.graphdb.AtlasUniqueKeyHandler;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.repository.graphdb.GraphIndexQueryParameters;
import org.apache.atlas.repository.graphdb.GremlinVersion;
import org.apache.atlas.repository.graphdb.janus.query.AtlasJanusGraphQuery;
import org.apache.atlas.repository.graphdb.utils.IteratorToIterableAdapter;
import org.apache.atlas.type.AtlasType;
import org.apache.commons.configuration.Configuration;
import org.apache.hadoop.util.StringUtils;
import org.apache.tinkerpop.gremlin.groovy.jsr223.GremlinGroovyScriptEngine;
import org.apache.tinkerpop.gremlin.jsr223.DefaultImportCustomizer;
import org.apache.tinkerpop.gremlin.process.traversal.P;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.__;
import org.apache.tinkerpop.gremlin.process.traversal.step.map.GraphStep;
import org.apache.tinkerpop.gremlin.process.traversal.step.util.ImmutablePath;
import org.apache.tinkerpop.gremlin.structure.Edge;
import org.apache.tinkerpop.gremlin.structure.Element;
import org.apache.tinkerpop.gremlin.structure.Property;
import org.apache.tinkerpop.gremlin.structure.Vertex;
import org.apache.tinkerpop.gremlin.structure.io.IoCore;
import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONMapper;
import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONWriter;
import org.janusgraph.core.Cardinality;
import org.janusgraph.core.JanusGraph;
import org.janusgraph.core.JanusGraphFactory;
import org.janusgraph.core.JanusGraphIndexQuery;
import org.janusgraph.core.PropertyKey;
import org.janusgraph.core.SchemaViolationException;
import org.janusgraph.core.schema.JanusGraphIndex;
import org.janusgraph.core.schema.JanusGraphManagement;
import org.janusgraph.core.schema.Parameter;
import org.janusgraph.diskstorage.BackendException;
import org.janusgraph.graphdb.database.StandardJanusGraph;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.script.Bindings;
import javax.script.ScriptEngine;
import javax.script.ScriptException;

import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import static org.apache.atlas.repository.Constants.INDEX_SEARCH_VERTEX_PREFIX_DEFAULT;
import static org.apache.atlas.repository.Constants.INDEX_SEARCH_VERTEX_PREFIX_PROPERTY;
import static org.apache.atlas.repository.graphdb.janus.AtlasJanusGraphDatabase.getGraphInstance;
import static org.apache.atlas.type.Constants.STATE_PROPERTY_KEY;

/**
 * Janus implementation of AtlasGraph.
 */
public class AtlasJanusGraph implements AtlasGraph<AtlasJanusVertex, AtlasJanusEdge> {
    private static final Logger LOG = LoggerFactory.getLogger(AtlasJanusGraph.class);

    private static final Parameter<?>[] EMPTY_PARAMETER_ARRAY = new Parameter[0];

    private static Configuration applicationProperties;

    private final ConvertGremlinValueFunction gremlinValueConversionFunction = new ConvertGremlinValueFunction();
    private final Set<String>                 multiProperties                = new HashSet<>();
    private final StandardJanusGraph          janusGraph;
    private final AtlasUniqueKeyHandler       uniqueKeyHandler;

    private final ThreadLocal<GremlinGroovyScriptEngine> scriptEngine = ThreadLocal.withInitial(() -> {
        DefaultImportCustomizer.Builder builder = DefaultImportCustomizer.build()
                .addClassImports(java.util.function.Function.class)
                .addMethodImports(__.class.getMethods())
                .addMethodImports(P.class.getMethods());
        return new GremlinGroovyScriptEngine(builder.create());
    });

    public AtlasJanusGraph() {
        this(getGraphInstance());
    }

    public AtlasJanusGraph(JanusGraph graphInstance) {
        //determine multi-properties once at startup
        JanusGraphManagement mgmt = null;

        try {
            mgmt = graphInstance.openManagement();

            Iterable<PropertyKey> keys = mgmt.getRelationTypes(PropertyKey.class);

            for (PropertyKey key : keys) {
                if (key.cardinality() != Cardinality.SINGLE) {
                    multiProperties.add(key.name());
                }
            }

            if (StringUtils.equalsIgnoreCase(AtlasConfiguration.STORAGE_BACKEND_TYPE.getString(), "rdbms")) {
                uniqueKeyHandler = new AtlasJanusRdbmsUniqueKeyHandler();
            } else {
                uniqueKeyHandler = null;
            }
        } finally {
            if (mgmt != null) {
                mgmt.rollback();
            }
        }

        janusGraph = (StandardJanusGraph) graphInstance;
    }

    @Override
    public AtlasEdge<AtlasJanusVertex, AtlasJanusEdge> addEdge(AtlasVertex<AtlasJanusVertex, AtlasJanusEdge> outVertex, AtlasVertex<AtlasJanusVertex, AtlasJanusEdge> inVertex, String edgeLabel) {
        try {
            Vertex oV   = outVertex.getV().getWrappedElement();
            Vertex iV   = inVertex.getV().getWrappedElement();
            Edge   edge = oV.addEdge(edgeLabel, iV);

            return GraphDbObjectFactory.createEdge(this, edge);
        } catch (SchemaViolationException e) {
            throw new AtlasSchemaViolationException(e);
        }
    }

    @Override
    public AtlasEdge<AtlasJanusVertex, AtlasJanusEdge> getEdgeBetweenVertices(AtlasVertex<AtlasJanusVertex, AtlasJanusEdge> fromVertex, AtlasVertex<AtlasJanusVertex, AtlasJanusEdge> toVertex, String edgeLabel) {
        GraphTraversal<?, ?> gt = V(fromVertex.getId()).outE(edgeLabel).where(__.otherV().hasId(toVertex.getId()));

        Edge gremlinEdge = getFirstActiveEdge(gt);

        return (gremlinEdge != null) ? GraphDbObjectFactory.createEdge(this, gremlinEdge) : null;
    }

    @Override
    public AtlasVertex<AtlasJanusVertex, AtlasJanusEdge> addVertex() {
        Vertex result = getGraph().addVertex();

        return GraphDbObjectFactory.createVertex(this, result);
    }

    @Override
    public void removeEdge(AtlasEdge<AtlasJanusVertex, AtlasJanusEdge> edge) {
        Edge wrapped = edge.getE().getWrappedElement();

        wrapped.remove();
    }

    @Override
    public void removeVertex(AtlasVertex<AtlasJanusVertex, AtlasJanusEdge> vertex) {
        Vertex wrapped = vertex.getV().getWrappedElement();

        wrapped.remove();
    }

    @Override
    public AtlasEdge<AtlasJanusVertex, AtlasJanusEdge> getEdge(String edgeId) {
        Iterator<Edge> it = getGraph().edges(edgeId);
        Edge           e  = getSingleElement(it, edgeId);

        return GraphDbObjectFactory.createEdge(this, e);
    }

    @Override
    public Iterable<AtlasEdge<AtlasJanusVertex, AtlasJanusEdge>> getEdges() {
        Iterator<Edge> edges = getGraph().edges();

        return wrapEdges(edges);
    }

    @Override
    public Iterable<AtlasVertex<AtlasJanusVertex, AtlasJanusEdge>> getVertices() {
        Iterator<Vertex> vertices = getGraph().vertices();

        return wrapVertices(vertices);
    }

    @Override
    public AtlasVertex<AtlasJanusVertex, AtlasJanusEdge> getVertex(String vertexId) {
        Iterator<Vertex> it     = getGraph().vertices(vertexId);
        Vertex           vertex = getSingleElement(it, vertexId);

        return GraphDbObjectFactory.createVertex(this, vertex);
    }

    @Override
    public Set<String> getEdgeIndexKeys() {
        return getIndexKeys(Edge.class);
    }

    @Override
    public Set<String> getVertexIndexKeys() {
        return getIndexKeys(Vertex.class);
    }

    @Override
    public Iterable<AtlasVertex<AtlasJanusVertex, AtlasJanusEdge>> getVertices(String key, Object value) {
        AtlasGraphQuery<AtlasJanusVertex, AtlasJanusEdge> query = query();

        query.has(key, value);

        return query.vertices();
    }

    @Override
    public AtlasGraphQuery<AtlasJanusVertex, AtlasJanusEdge> query() {
        return new AtlasJanusGraphQuery(this);
    }

    @Override
    public AtlasGraphTraversal<AtlasVertex<?, ?>, AtlasEdge<?, ?>> V(Object... vertexIds) {
        AtlasJanusGraphTraversal traversal = new AtlasJanusGraphTraversal(this, getGraph().traversal());

        traversal.getBytecode().addStep(GraphTraversal.Symbols.V, vertexIds);
        traversal.addStep(new GraphStep<>(traversal, Vertex.class, true, vertexIds));

        return (AtlasGraphTraversal) traversal;
    }

    @Override
    public AtlasGraphTraversal<AtlasVertex<?, ?>, AtlasEdge<?, ?>> E(Object... edgeIds) {
        AtlasJanusGraphTraversal traversal = new AtlasJanusGraphTraversal(this, getGraph().traversal());

        traversal.getBytecode().addStep(GraphTraversal.Symbols.E, edgeIds);
        traversal.addStep(new GraphStep<>(traversal, Vertex.class, true, edgeIds));

        return (AtlasGraphTraversal) traversal;
    }

    @Override
    public AtlasIndexQuery<AtlasJanusVertex, AtlasJanusEdge> indexQuery(String indexName, String graphQuery) {
        return indexQuery(indexName, graphQuery, 0, null);
    }

    @Override
    public AtlasIndexQuery<AtlasJanusVertex, AtlasJanusEdge> indexQuery(String indexName, String graphQuery, int offset) {
        return indexQuery(indexName, graphQuery, offset, null);
    }

    /**
     * Creates an index query.
     *
     * @param indexQueryParameters the parameterObject containing the information needed for creating the index.
     */
    public AtlasIndexQuery<AtlasJanusVertex, AtlasJanusEdge> indexQuery(GraphIndexQueryParameters indexQueryParameters) {
        return indexQuery(indexQueryParameters.getIndexName(), indexQueryParameters.getGraphQueryString(), indexQueryParameters.getOffset(), indexQueryParameters.getIndexQueryParameters());
    }

    @Override
    public AtlasGraphManagement getManagementSystem() {
        return new AtlasJanusGraphManagement(this, getGraph().openManagement());
    }

    @Override
    public void commit() {
        getGraph().tx().commit();
    }

    @Override
    public void rollback() {
        getGraph().tx().rollback();
    }

    @Override
    public void shutdown() {
        getGraph().close();
    }

    @Override
    public void clear() {
        JanusGraph graph = getGraph();

        if (graph.isOpen()) {
            // only a shut down graph can be cleared
            graph.close();
        }

        try {
            JanusGraphFactory.drop(graph);
        } catch (BackendException ignoreEx) {
        }
    }

    @Override
    public Set<?> getOpenTransactions() {
        return janusGraph.getOpenTransactions();
    }

    @Override
    public void exportToGson(OutputStream os) throws IOException {
        GraphSONMapper         mapper  = getGraph().io(IoCore.graphson()).mapper().create();
        GraphSONWriter.Builder builder = GraphSONWriter.build();

        builder.mapper(mapper);

        GraphSONWriter writer = builder.create();

        writer.writeGraph(os, getGraph());
    }

    @Override
    public GroovyExpression generatePersisentToLogicalConversionExpression(GroovyExpression expr, AtlasType type) {
        //nothing special needed, value is stored in required type
        return expr;
    }

    @Override
    public boolean isPropertyValueConversionNeeded(AtlasType type) {
        return false;
    }

    @Override
    public GremlinVersion getSupportedGremlinVersion() {
        return GremlinVersion.THREE;
    }

    @Override
    public boolean requiresInitialIndexedPredicate() {
        return false;
    }

    @Override
    public GroovyExpression getInitialIndexedPredicate(GroovyExpression parent) {
        return parent;
    }

    @Override
    public GroovyExpression addOutputTransformationPredicate(GroovyExpression expr, boolean isSelect, boolean isPath) {
        return expr;
    }

    @Override
    public GremlinGroovyScriptEngine getGremlinScriptEngine() {
        return scriptEngine.get();
    }

    @Override
    public void releaseGremlinScriptEngine(ScriptEngine scriptEngine) {
        if (scriptEngine instanceof GremlinGroovyScriptEngine) {
            try {
                ((GremlinGroovyScriptEngine) scriptEngine).reset();
            } catch (Exception e) {
                // ignore
            }
        }
    }

    @Override
    public Object executeGremlinScript(String query, boolean isPath) throws AtlasBaseException {
        Object result = executeGremlinScript(query);

        return convertGremlinValue(result);
    }

    @Override
    public Object executeGremlinScript(ScriptEngine scriptEngine, Map<? extends String, ?> userBindings, String query, boolean isPath) throws ScriptException {
        Bindings bindings = scriptEngine.createBindings();

        bindings.putAll(userBindings);
        bindings.put("g", getGraph().traversal());

        Object result = scriptEngine.eval(query, bindings);

        return convertGremlinValue(result);
    }

    @Override
    public boolean isMultiProperty(String propertyName) {
        return multiProperties.contains(propertyName);
    }

    @Override
    public AtlasIndexQueryParameter indexQueryParameter(String parameterName, String parameterValue) {
        return new AtlasJanusIndexQueryParameter(parameterName, parameterValue);
    }

    @Override
    public AtlasGraphIndexClient getGraphIndexClient() throws AtlasException {
        try {
            initApplicationProperties();

            return new AtlasJanusGraphIndexClient(applicationProperties);
        } catch (Exception e) {
            LOG.error("Error encountered in creating Graph Index Client.", e);
            throw new AtlasException(e);
        }
    }

    public JanusGraph getGraph() {
        return this.janusGraph;
    }

    @Override
    public List<AtlasVertex> getAllEdgesVertices(AtlasVertex vertex) {
        GraphTraversal gt = V(vertex.getId()).both();
        List<AtlasVertex> resultList = new ArrayList<>();
        while (gt.hasNext()) {
            Vertex v = (Vertex) gt.next();
            resultList.add(GraphDbObjectFactory.createVertex(this, v));
        }
        return resultList;
    }

    @Override
    public AtlasUniqueKeyHandler getUniqueKeyHandler() {
        return uniqueKeyHandler;
    }

    public Iterable<AtlasVertex<AtlasJanusVertex, AtlasJanusEdge>> wrapVertices(Iterable<? extends Vertex> it) {
        return Iterables.transform(it, (Function<Vertex, AtlasVertex<AtlasJanusVertex, AtlasJanusEdge>>) input -> GraphDbObjectFactory.createVertex(AtlasJanusGraph.this, input));
    }

    public Iterable<AtlasEdge<AtlasJanusVertex, AtlasJanusEdge>> wrapEdges(Iterator<? extends Edge> it) {
        Iterable<? extends Edge> iterable = new IteratorToIterableAdapter<>(it);

        return wrapEdges(iterable);
    }

    public Iterable<AtlasEdge<AtlasJanusVertex, AtlasJanusEdge>> wrapEdges(Iterable<? extends Edge> it) {
        return Iterables.transform(it, (Function<Edge, AtlasEdge<AtlasJanusVertex, AtlasJanusEdge>>) input -> GraphDbObjectFactory.createEdge(AtlasJanusGraph.this, input));
    }

    public void addMultiProperties(Set<String> names) {
        multiProperties.addAll(names);
    }

    String getIndexFieldName(AtlasPropertyKey propertyKey, JanusGraphIndex graphIndex, Parameter<?>... parameters) {
        PropertyKey janusKey = AtlasJanusObjectFactory.createPropertyKey(propertyKey);

        if (parameters == null) {
            parameters = EMPTY_PARAMETER_ARRAY;
        }

        return janusGraph.getIndexSerializer().getDefaultFieldName(janusKey, parameters, graphIndex.getBackingIndex());
    }

    private AtlasIndexQuery<AtlasJanusVertex, AtlasJanusEdge> indexQuery(String indexName, String graphQuery, int offset, List<AtlasIndexQueryParameter> indexQueryParameterList) {
        String               prefix = getIndexQueryPrefix();
        JanusGraphIndexQuery query  = getGraph().indexQuery(indexName, graphQuery).setElementIdentifier(prefix).offset(offset);

        if (indexQueryParameterList != null && !indexQueryParameterList.isEmpty()) {
            for (AtlasIndexQueryParameter indexQueryParameter : indexQueryParameterList) {
                query = query.addParameter(new Parameter<>(indexQueryParameter.getParameterName(), indexQueryParameter.getParameterValue()));
            }
        }

        return new AtlasJanusIndexQuery(this, query);
    }

    private String getIndexQueryPrefix() {
        final String ret;

        initApplicationProperties();

        if (applicationProperties == null) {
            ret = INDEX_SEARCH_VERTEX_PREFIX_DEFAULT;
        } else {
            ret = applicationProperties.getString(INDEX_SEARCH_VERTEX_PREFIX_PROPERTY, INDEX_SEARCH_VERTEX_PREFIX_DEFAULT);
        }

        return ret;
    }

    private Iterable<AtlasVertex<AtlasJanusVertex, AtlasJanusEdge>> wrapVertices(Iterator<? extends Vertex> it) {
        Iterable<? extends Vertex> iterable = new IteratorToIterableAdapter<>(it);

        return wrapVertices(iterable);
    }

    private static <T> T getSingleElement(Iterator<T> it, String id) {
        if (!it.hasNext()) {
            return null;
        }

        T element = it.next();

        if (it.hasNext()) {
            throw new RuntimeException("Multiple items were found with the id " + id);
        }

        return element;
    }

    private Object convertGremlinValue(Object rawValue) {
        if (rawValue instanceof Vertex) {
            return GraphDbObjectFactory.createVertex(this, (Vertex) rawValue);
        } else if (rawValue instanceof Edge) {
            return GraphDbObjectFactory.createEdge(this, (Edge) rawValue);
        } else if (rawValue instanceof Map) {
            Map<String, Object> rowValue = (Map<String, Object>) rawValue;

            return Maps.transformValues(rowValue, gremlinValueConversionFunction);
        } else if (rawValue instanceof ImmutablePath) {
            ImmutablePath path = (ImmutablePath) rawValue;

            return convertGremlinValue(path.objects());
        } else if (rawValue instanceof List) {
            return Lists.transform((List<?>) rawValue, gremlinValueConversionFunction);
        } else if (rawValue instanceof Collection) {
            throw new UnsupportedOperationException("Unhandled collection type: " + rawValue.getClass());
        }

        return rawValue;
    }

    private Set<String> getIndexKeys(Class<? extends Element> janusGraphElementClass) {
        Set<String>          result = new HashSet<>();
        JanusGraphManagement mgmt   = null;

        try {
            mgmt = getGraph().openManagement();

            Iterable<JanusGraphIndex> indices = mgmt.getGraphIndexes(janusGraphElementClass);

            for (JanusGraphIndex index : indices) {
                result.add(index.name());
            }
        } finally {
            if (mgmt != null) {
                mgmt.commit();
            }
        }

        return result;
    }

    private Object executeGremlinScript(String gremlinQuery) throws AtlasBaseException {
        GremlinGroovyScriptEngine scriptEngine = getGremlinScriptEngine();

        try {
            Bindings bindings = scriptEngine.createBindings();

            bindings.put("graph", getGraph());
            bindings.put("g", getGraph().traversal());

            return scriptEngine.eval(gremlinQuery, bindings);
        } catch (ScriptException e) {
            throw new AtlasBaseException(AtlasErrorCode.GREMLIN_SCRIPT_EXECUTION_FAILED, e, gremlinQuery);
        } finally {
            releaseGremlinScriptEngine(scriptEngine);
        }
    }

    private void initApplicationProperties() {
        if (applicationProperties == null) {
            try {
                applicationProperties = ApplicationProperties.get();
            } catch (AtlasException ex) {
                // ignore
            }
        }
    }

    private Edge getFirstActiveEdge(GraphTraversal<?, ?> gt) {
        if (gt != null) {
            while (gt.hasNext()) {
                Edge gremlinEdge = (Edge) gt.next();

                if (gremlinEdge != null) {
                    Property<?> property = gremlinEdge.property(STATE_PROPERTY_KEY);

                    if (property.isPresent() && property.value().equals(AtlasEntity.Status.ACTIVE.toString())) {
                        return gremlinEdge;
                    }
                }
            }
        }

        return null;
    }

    private final class ConvertGremlinValueFunction implements Function<Object, Object> {
        @Override
        public Object apply(Object input) {
            return convertGremlinValue(input);
        }
    }
}
