/**
 * 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 org.apache.atlas.AtlasException;
import org.apache.atlas.graph.GraphSandboxUtil;
import org.apache.atlas.repository.Constants;
import org.apache.atlas.repository.graphdb.AtlasCardinality;
import org.apache.atlas.repository.graphdb.AtlasEdge;
import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasGraphManagement;
import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
import org.apache.atlas.repository.graphdb.AtlasGraphQuery.ComparisionOperator;
import org.apache.atlas.repository.graphdb.AtlasPropertyKey;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.runner.LocalSolrRunner;
import org.apache.atlas.typesystem.types.DataTypes.TypeCategory;
import org.apache.commons.configuration.Configuration;
import org.testng.annotations.AfterClass;
import org.testng.annotations.Test;

import java.math.BigDecimal;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import static org.apache.atlas.graph.GraphSandboxUtil.useLocalSolr;
import static org.apache.atlas.repository.graphdb.janus.AtlasJanusGraphDatabase.initJanusGraph;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

/**
 * Sanity test of basic graph operations using the Janus graphdb
 * abstraction layer implementation.
 */
public class AtlasJanusDatabaseTest {
    private AtlasGraph<?, ?> atlasGraph;

    @AfterClass
    public void cleanup() throws Exception {
        if (atlasGraph != null) {
            atlasGraph.clear();
            atlasGraph = null;
        }

        if (useLocalSolr()) {
            LocalSolrRunner.stop();
        }
    }

    @Test
    public void initializationFailureShouldThrowRuntimeException() throws AtlasException {
        Configuration       cfg       = AtlasJanusGraphDatabase.getConfiguration();
        Map<String, Object> cfgBackup = new HashMap<>();
        Iterator<String>    keys      = cfg.getKeys();

        while (keys.hasNext()) {
            String key = keys.next();

            cfgBackup.put(key, cfg.getProperty(key));
        }

        try {
            cfg.clear();

            initJanusGraph(cfg);

            fail("Should have thrown an exception!");
        } catch (RuntimeException ex) {
            assertTrue(true);

            for (Map.Entry<String, Object> entry : cfgBackup.entrySet()) {
                cfg.setProperty(entry.getKey(), entry.getValue());
            }
        }
    }

    @Test
    public <V, E> void testPropertyDataTypes() throws Exception {
        // primitives
        AtlasGraph<V, E> graph = getGraph();

        testProperty(graph, "booleanProperty", Boolean.TRUE);
        testProperty(graph, "booleanProperty", Boolean.FALSE);
        testProperty(graph, "booleanProperty", Boolean.TRUE);
        testProperty(graph, "booleanProperty", Boolean.FALSE);

        testProperty(graph, "byteProperty", Byte.MAX_VALUE);
        testProperty(graph, "byteProperty", Byte.MIN_VALUE);
        testProperty(graph, "byteProperty", Byte.MAX_VALUE);
        testProperty(graph, "byteProperty", Byte.MIN_VALUE);

        testProperty(graph, "shortProperty", Short.MAX_VALUE);
        testProperty(graph, "shortProperty", Short.MIN_VALUE);
        testProperty(graph, "shortProperty", Short.MAX_VALUE);
        testProperty(graph, "shortProperty", Short.MIN_VALUE);

        testProperty(graph, "intProperty", Integer.MAX_VALUE);
        testProperty(graph, "intProperty", Integer.MIN_VALUE);
        testProperty(graph, "intProperty", Integer.MAX_VALUE);
        testProperty(graph, "intProperty", Integer.MIN_VALUE);

        testProperty(graph, "longProperty", Long.MIN_VALUE);
        testProperty(graph, "longProperty", Long.MAX_VALUE);
        testProperty(graph, "longProperty", Long.MIN_VALUE);
        testProperty(graph, "longProperty", Long.MAX_VALUE);

        testProperty(graph, "doubleProperty", Double.MAX_VALUE);
        testProperty(graph, "doubleProperty", Double.MIN_VALUE);
        testProperty(graph, "doubleProperty", Double.MAX_VALUE);
        testProperty(graph, "doubleProperty", Double.MIN_VALUE);

        testProperty(graph, "floatProperty", Float.MAX_VALUE);
        testProperty(graph, "floatProperty", Float.MIN_VALUE);
        testProperty(graph, "floatProperty", Float.MAX_VALUE);
        testProperty(graph, "floatProperty", Float.MIN_VALUE);

        // enumerations - TypeCategory
        testProperty(graph, "typeCategoryProperty", TypeCategory.CLASS);

        // biginteger
        testProperty(graph, "bigIntegerProperty", new BigInteger(String.valueOf(Long.MAX_VALUE)).multiply(BigInteger.TEN));

        // bigdecimal
        BigDecimal bigDecimal = new BigDecimal(Double.MAX_VALUE);

        testProperty(graph, "bigDecimalProperty", bigDecimal.multiply(bigDecimal));
    }

    @Test
    public <V, E> void testMultiplicityOnePropertySupport() throws Exception {
        AtlasGraph<V, E>  graph  = getGraph();
        AtlasVertex<V, E> vertex = graph.addVertex();

        vertex.setProperty("name", "Jeff");
        vertex.setProperty("location", "Littleton");

        assertEquals(vertex.getProperty("name", String.class), "Jeff");
        assertEquals(vertex.getProperty("location", String.class), "Littleton");

        AtlasVertex<V, E> vertexCopy = graph.getVertex(vertex.getId().toString());

        assertEquals(vertexCopy.getProperty("name", String.class), "Jeff");
        assertEquals(vertexCopy.getProperty("location", String.class), "Littleton");

        assertTrue(vertexCopy.getPropertyKeys().contains("name"));
        assertTrue(vertexCopy.getPropertyKeys().contains("location"));

        assertTrue(vertexCopy.getPropertyValues("name", String.class).contains("Jeff"));
        assertTrue(vertexCopy.getPropertyValues("location", String.class).contains("Littleton"));
        assertTrue(vertexCopy.getPropertyValues("test", String.class).isEmpty());
        assertNull(vertexCopy.getProperty("test", String.class));

        vertex.removeProperty("name");
        assertFalse(vertex.getPropertyKeys().contains("name"));
        assertNull(vertex.getProperty("name", String.class));
        assertTrue(vertex.getPropertyValues("name", String.class).isEmpty());

        vertexCopy = graph.getVertex(vertex.getId().toString());
        assertFalse(vertexCopy.getPropertyKeys().contains("name"));
        assertNull(vertexCopy.getProperty("name", String.class));
        assertTrue(vertexCopy.getPropertyValues("name", String.class).isEmpty());
    }

    @Test
    public <V, E> void testRemoveEdge() throws Exception {
        AtlasGraph<V, E>  graph = getGraph();
        AtlasVertex<V, E> v1    = graph.addVertex();
        AtlasVertex<V, E> v2    = graph.addVertex();
        AtlasEdge<V, E>   edge  = graph.addEdge(v1, v2, "knows");

        // make sure the edge exists
        AtlasEdge<V, E> edgeCopy = graph.getEdge(edge.getId().toString());
        assertNotNull(edgeCopy);
        assertEquals(edgeCopy, edge);

        graph.removeEdge(edge);

        edgeCopy = graph.getEdge(edge.getId().toString());
        // should return null now, since edge was deleted
        assertNull(edgeCopy);
    }

    @Test
    public <V, E> void testRemoveVertex() throws Exception {
        AtlasGraph<V, E>  graph = getGraph();
        AtlasVertex<V, E> v1    = graph.addVertex();

        assertNotNull(graph.getVertex(v1.getId().toString()));

        graph.removeVertex(v1);

        assertNull(graph.getVertex(v1.getId().toString()));
    }

    @Test
    public <V, E> void testGetEdges() throws Exception {
        AtlasGraph<V, E>  graph = getGraph();
        AtlasVertex<V, E> v1    = graph.addVertex();
        AtlasVertex<V, E> v2    = graph.addVertex();
        AtlasVertex<V, E> v3    = graph.addVertex();

        AtlasEdge<V, E> knows  = graph.addEdge(v2, v1, "knows");
        AtlasEdge<V, E> eats   = graph.addEdge(v3, v1, "eats");
        AtlasEdge<V, E> drives = graph.addEdge(v3, v2, "drives");
        AtlasEdge<V, E> sleeps = graph.addEdge(v2, v3, "sleeps");

        assertEdgesMatch(v1.getEdges(AtlasEdgeDirection.IN), knows, eats);
        assertEdgesMatch(v1.getEdges(AtlasEdgeDirection.OUT));
        assertEdgesMatch(v1.getEdges(AtlasEdgeDirection.BOTH), knows, eats);

        assertEdgesMatch(v1.getEdges(AtlasEdgeDirection.IN, "knows"), knows);
        assertEdgesMatch(v1.getEdges(AtlasEdgeDirection.OUT, "knows"));
        assertEdgesMatch(v1.getEdges(AtlasEdgeDirection.BOTH, "knows"), knows);

        assertEdgesMatch(v2.getEdges(AtlasEdgeDirection.IN), drives);
        assertEdgesMatch(v2.getEdges(AtlasEdgeDirection.OUT), knows, sleeps);
        assertEdgesMatch(v2.getEdges(AtlasEdgeDirection.BOTH), knows, sleeps, drives);

        assertEdgesMatch(v2.getEdges(AtlasEdgeDirection.BOTH, "delivers"));
    }

    @Test
    public <V, E> void testMultiplictyManyPropertySupport() throws Exception {
        AtlasGraph<V, E>  graph    = getGraph();
        AtlasVertex<V, E> vertex   = graph.addVertex();
        String            vertexId = vertex.getId().toString();

        vertex.setProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait1");
        vertex.setProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait2");

        assertEquals(vertex.getPropertyValues(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class).size(), 2);

        vertex.addProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait3");
        vertex.addProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait4");

        assertTrue(vertex.getPropertyKeys().contains(Constants.TRAIT_NAMES_PROPERTY_KEY));
        validateMultManyPropertiesInVertex(vertex);
        // fetch a copy of the vertex, make sure result
        // is the same

        validateMultManyPropertiesInVertex(graph.getVertex(vertexId));
    }

    @Test
    public <V, E> void testListProperties() throws Exception {
        AtlasGraph<V, E>  graph       = getGraph();
        AtlasVertex<V, E> vertex      = graph.addVertex();
        List<String>      colorsToSet = new ArrayList<>();

        colorsToSet.add("red");
        colorsToSet.add("blue");
        colorsToSet.add("green");

        vertex.setListProperty("colors", colorsToSet);

        List<String> colors = vertex.getListProperty("colors");

        assertTrue(colors.contains("red"));
        assertTrue(colors.contains("blue"));
        assertTrue(colors.contains("green"));

        AtlasVertex<V, E> vertexCopy = graph.getVertex(vertex.getId().toString());

        colors = vertexCopy.getListProperty("colors");

        assertTrue(colors.contains("red"));
        assertTrue(colors.contains("blue"));
        assertTrue(colors.contains("green"));
    }

    @Test
    public <V, E> void testRemoveProperty() throws Exception {
        AtlasGraph<V, E>  graph  = getGraph();
        AtlasVertex<V, E> vertex = graph.addVertex();

        vertex.addProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait1");
        vertex.addProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait1");
        vertex.setProperty("name", "Jeff");

        // remove existing property - multiplicity one
        vertex.removeProperty("jeff");

        assertFalse(vertex.getPropertyKeys().contains("jeff"));

        // remove existing property - multiplicity many
        vertex.removeProperty(Constants.TRAIT_NAMES_PROPERTY_KEY);
        assertFalse(vertex.getPropertyKeys().contains(Constants.TRAIT_NAMES_PROPERTY_KEY));

        AtlasVertex<V, E> vertexCopy = graph.getVertex(vertex.getId().toString());

        assertFalse(vertexCopy.getPropertyKeys().contains("jeff"));
        assertFalse(vertexCopy.getPropertyKeys().contains(Constants.TRAIT_NAMES_PROPERTY_KEY));

        // remove non-existing property
        vertex.removeProperty(Constants.TRAIT_NAMES_PROPERTY_KEY);
        vertex.removeProperty("jeff");
    }

    @Test
    public <V, E> void getGetGraphQueryForVertices() throws Exception {
        AtlasGraph<V, E>  graph = getGraph();
        AtlasVertex<V, E> v1    = graph.addVertex();
        AtlasVertex<V, E> v2    = graph.addVertex();
        AtlasVertex<V, E> v3    = graph.addVertex();

        v1.setProperty("name", "Jeff");
        v1.setProperty("weight", 1);

        v2.setProperty("name", "Fred");
        v2.setProperty("weight", 2);

        v3.setProperty("name", "Chris");
        v3.setProperty("weight", 3);

        AtlasEdge<V, E> knows = graph.addEdge(v2, v1, "knows");

        knows.setProperty("weight", 1);

        AtlasEdge<V, E> eats = graph.addEdge(v3, v1, "eats");

        eats.setProperty("weight", 2);

        AtlasEdge<V, E> drives = graph.addEdge(v3, v2, "drives");

        drives.setProperty("weight", 3);

        AtlasEdge<V, E> sleeps = graph.addEdge(v2, v3, "sleeps");

        sleeps.setProperty("weight", 4);

        testExecuteGraphQuery("name", null, "Jeff", v1);
        testExecuteGraphQuery("weight", ComparisionOperator.EQUAL, 2, v2);
        testExecuteGraphQuery("weight", ComparisionOperator.GREATER_THAN_EQUAL, 2, v2, v3);
        testExecuteGraphQuery("weight", ComparisionOperator.LESS_THAN_EQUAL, 2, v2, v1);
    }

    @Test
    public <V, E> void testAddMultManyPropertyValueTwice() throws Exception {
        AtlasGraph<V, E>  graph    = getGraph();
        AtlasVertex<V, E> vertex   = graph.addVertex();
        String            vertexId = vertex.getId().toString();

        vertex.setProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait1");
        vertex.setProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait1");
        vertex.addProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait2");
        vertex.addProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, "trait2");

        validateDuplicatePropertyVertex(vertex);

        // fetch a copy of the vertex, make sure result is the same

        validateDuplicatePropertyVertex(graph.getVertex(vertexId));
    }

    private <V, E> AtlasGraph<V, E> getGraph() throws Exception {
        GraphSandboxUtil.create();

        if (useLocalSolr()) {
            LocalSolrRunner.start();
        }

        if (atlasGraph == null) {
            AtlasJanusGraphDatabase db = new AtlasJanusGraphDatabase();

            atlasGraph = db.getGraph();

            try (AtlasGraphManagement mgmt = atlasGraph.getManagementSystem()) {
                // create the index (which defines these properties as being mult
                // many)
                for (String propertyName : new String[] {"__superTypeNames", "__traitNames"}) {
                    AtlasPropertyKey propertyKey = mgmt.getPropertyKey(propertyName);

                    if (propertyKey == null) {
                        propertyKey = mgmt.makePropertyKey(propertyName, String.class, AtlasCardinality.SET);

                        mgmt.createVertexCompositeIndex(propertyName, false, Collections.singletonList(propertyKey));
                    }
                }

                mgmt.setIsSuccess(true);
            }
        }

        return (AtlasGraph<V, E>) atlasGraph;
    }

    private <V, E> void testProperty(AtlasGraph<V, E> graph, String name, Object value) {
        AtlasVertex<V, E> vertex = graph.addVertex();

        vertex.setProperty(name, value);

        assertEquals(value, vertex.getProperty(name, value.getClass()));

        AtlasVertex<V, E> loaded = graph.getVertex(vertex.getId().toString());

        assertEquals(value, loaded.getProperty(name, value.getClass()));
    }

    @SafeVarargs
    private final <V, E> void assertEdgesMatch(Iterable<AtlasEdge<V, E>> edgesIt, AtlasEdge<V, E>... expected) {
        List<AtlasEdge<V, E>> edges = toList(edgesIt);

        assertEquals(expected.length, edges.size());

        for (AtlasEdge<V, E> edge : expected) {
            assertTrue(edges.contains(edge));
        }
    }

    private <V, E> void validateMultManyPropertiesInVertex(AtlasVertex<V, E> vertex) {
        assertTrue(vertex.getPropertyKeys().contains(Constants.TRAIT_NAMES_PROPERTY_KEY));

        Collection<String> traitNames = vertex.getPropertyValues(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class);

        assertTrue(traitNames.contains("trait1"));
        assertTrue(traitNames.contains("trait2"));
        assertTrue(traitNames.contains("trait3"));
        assertTrue(traitNames.contains("trait4"));

        try {
            vertex.getProperty(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class);

            fail("Expected exception not thrown");
        } catch (IllegalStateException expected) {
            // multiple property values exist
        }
    }

    @SafeVarargs
    private final <V, E> void testExecuteGraphQuery(String property, ComparisionOperator op, Object value, AtlasVertex<V, E>... expected) throws Exception {
        AtlasGraph<V, E>      graph = getGraph();
        AtlasGraphQuery<V, E> query = graph.query();

        if (op != null) {
            query.has(property, op, value);
        } else {
            query.has(property, value);
        }

        Iterable<? extends AtlasVertex<V, E>> result = query.vertices();
        List<AtlasVertex<V, E>>               list   = toList(result);

        assertEquals(expected.length, list.size());

        for (AtlasVertex<V, E> vertex : expected) {
            assertTrue(list.contains(vertex));
        }
    }

    private <V, E> void validateDuplicatePropertyVertex(AtlasVertex<V, E> vertex) {
        assertEquals(vertex.getPropertyValues(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class).size(), 2);
        assertTrue(vertex.getPropertyKeys().contains(Constants.TRAIT_NAMES_PROPERTY_KEY));

        Collection<String> traitNames = vertex.getPropertyValues(Constants.TRAIT_NAMES_PROPERTY_KEY, String.class);

        assertTrue(traitNames.contains("trait1"));
        assertTrue(traitNames.contains("trait2"));
    }

    private static <T> List<T> toList(Iterable<? extends T> iterable) {
        List<T> result = new ArrayList<>();

        for (T item : iterable) {
            result.add(item);
        }

        return result;
    }
}
