/**
 * 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.graph.GraphSandboxUtil;
import org.apache.atlas.repository.Constants;
import org.apache.atlas.repository.graphdb.AtlasCardinality;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.AtlasGraphManagement;
import org.apache.atlas.repository.graphdb.AtlasPropertyKey;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.runner.LocalSolrRunner;
import org.testng.annotations.AfterClass;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeClass;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import static org.apache.atlas.graph.GraphSandboxUtil.useLocalSolr;

public abstract class AbstractGraphDatabaseTest {
    protected static final String WEIGHT_PROPERTY    = "weight";
    protected static final String TRAIT_NAMES        = Constants.TRAIT_NAMES_PROPERTY_KEY;
    protected static final String TYPE_PROPERTY_NAME = "__type";
    protected static final String TYPESYSTEM         = "TYPESYSTEM";

    private static final String BACKING_INDEX_NAME = "backing";

    protected List<AtlasVertex<?, ?>> newVertices = new ArrayList<>();

    private AtlasGraph<?, ?> graph;

    @BeforeClass
    public static void createIndices() throws Exception {
        GraphSandboxUtil.create();

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

        AtlasJanusGraphDatabase db = new AtlasJanusGraphDatabase();

        try (AtlasGraphManagement mgmt = db.getGraph().getManagementSystem()) {
            if (mgmt.getGraphIndex(BACKING_INDEX_NAME) == null) {
                mgmt.createVertexMixedIndex(BACKING_INDEX_NAME, Constants.BACKING_INDEX, Collections.emptyList());
            }
            mgmt.makePropertyKey("age13", Integer.class, AtlasCardinality.SINGLE);

            createIndices(mgmt, "name", String.class, false, AtlasCardinality.SINGLE);
            createIndices(mgmt, WEIGHT_PROPERTY, Integer.class, false, AtlasCardinality.SINGLE);
            createIndices(mgmt, "size15", String.class, false, AtlasCardinality.SINGLE);
            createIndices(mgmt, "typeName", String.class, false, AtlasCardinality.SINGLE);
            createIndices(mgmt, "__type", String.class, false, AtlasCardinality.SINGLE);
            createIndices(mgmt, Constants.GUID_PROPERTY_KEY, String.class, true, AtlasCardinality.SINGLE);
            createIndices(mgmt, Constants.TRAIT_NAMES_PROPERTY_KEY, String.class, false, AtlasCardinality.SET);
            createIndices(mgmt, Constants.SUPER_TYPES_PROPERTY_KEY, String.class, false, AtlasCardinality.SET);

            mgmt.setIsSuccess(true);
        }
    }

    @AfterClass
    public static void cleanUp() throws Exception {
        AtlasJanusGraph graph = new AtlasJanusGraph();

        graph.clear();

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

    @AfterMethod
    public void commitGraph() {
        //force any pending actions to be committed so we can be sure they don't cause errors.
        pushChangesAndFlushCache();
        getGraph().commit();
    }

    @AfterMethod
    public void removeVertices() {
        for (AtlasVertex vertex : newVertices) {
            if (vertex.exists()) {
                getGraph().removeVertex(vertex);
            }
        }

        getGraph().commit();
        newVertices.clear();
    }

    protected <V, E> void pushChangesAndFlushCache() {
        getGraph().commit();
    }

    protected final <V, E> AtlasGraph<V, E> getGraph() {
        if (graph == null) {
            graph = new AtlasJanusGraph();
        }

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

    protected AtlasJanusGraph getAtlasJanusGraph() {
        AtlasGraph<?, ?> g = getGraph();

        return (AtlasJanusGraph) g;
    }

    protected final <V, E> AtlasVertex<V, E> createVertex(AtlasGraph<V, E> theGraph) {
        AtlasVertex<V, E> vertex = theGraph.addVertex();

        newVertices.add(vertex);

        return vertex;
    }

    protected void runSynchronouslyInNewThread(final Runnable r) throws Throwable {
        RunnableWrapper wrapper = new RunnableWrapper(r);
        Thread          th      = new Thread(wrapper);

        th.start();
        th.join();

        Throwable ex = wrapper.getExceptionThrown();

        if (ex != null) {
            throw ex;
        }
    }

    private static void createIndices(AtlasGraphManagement management, String propertyName, Class<?> propertyClass, boolean isUnique, AtlasCardinality cardinality) {
        if (management.containsPropertyKey(propertyName)) {
            //index was already created
            return;
        }

        AtlasPropertyKey key = management.makePropertyKey(propertyName, propertyClass, cardinality);

        try {
            if (propertyClass != Integer.class) {
                management.addMixedIndex(BACKING_INDEX_NAME, key, false);
            }
        } catch (Throwable t) {
            //ok
            t.printStackTrace();
        }

        try {
            management.createVertexCompositeIndex(propertyName, isUnique, Collections.singletonList(key));
        } catch (Throwable t) {
            //ok
            t.printStackTrace();
        }
    }

    private static final class RunnableWrapper implements Runnable {
        private final Runnable  r;
        private       Throwable exceptionThrown;

        private RunnableWrapper(Runnable r) {
            this.r = r;
        }

        @Override
        public void run() {
            try {
                r.run();
            } catch (Throwable e) {
                exceptionThrown = e;
            }
        }

        public Throwable getExceptionThrown() {
            return exceptionThrown;
        }
    }
}
