/**
 * 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.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableMap;
import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasConfiguration;
import org.apache.atlas.AtlasException;
import org.apache.atlas.repository.graphdb.AtlasGraph;
import org.apache.atlas.repository.graphdb.GraphDatabase;
import org.apache.atlas.repository.graphdb.janus.serializer.BigDecimalSerializer;
import org.apache.atlas.repository.graphdb.janus.serializer.BigIntegerSerializer;
import org.apache.atlas.repository.graphdb.janus.serializer.TypeCategorySerializer;
import org.apache.atlas.typesystem.types.DataTypes.TypeCategory;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.configuration.ConfigurationConverter;
import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONMapper;
import org.janusgraph.core.JanusGraph;
import org.janusgraph.core.JanusGraphException;
import org.janusgraph.core.JanusGraphFactory;
import org.janusgraph.core.schema.JanusGraphManagement;
import org.janusgraph.diskstorage.StandardIndexProvider;
import org.janusgraph.diskstorage.StandardStoreManager;
import org.janusgraph.diskstorage.es.ElasticSearch7Index;
import org.janusgraph.diskstorage.hbase.HBaseStoreManager;
import org.janusgraph.diskstorage.rdbms.RdbmsStoreManager;
import org.janusgraph.diskstorage.solr.Solr6Index;
import org.janusgraph.graphdb.database.serialize.attribute.SerializableSerializer;
import org.janusgraph.graphdb.tinkerpop.JanusGraphIoRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;

import static org.apache.atlas.ApplicationProperties.DEFAULT_INDEX_RECOVERY;
import static org.apache.atlas.ApplicationProperties.INDEX_RECOVERY_CONF;

/**
 * Default implementation for Graph Provider that doles out JanusGraph.
 */
public class AtlasJanusGraphDatabase implements GraphDatabase<AtlasJanusVertex, AtlasJanusEdge> {
    private static final Logger LOG = LoggerFactory.getLogger(AtlasJanusGraphDatabase.class);

    /**
     * Constant for the configuration property that indicates the prefix.
     */
    public static final String GRAPH_PREFIX              = "atlas.graph";
    public static final String INDEX_BACKEND_CONF        = "index.search.backend";
    public static final String SOLR_ZOOKEEPER_URL        = "atlas.graph.index.search.solr.zookeeper-url";
    public static final String SOLR_ZOOKEEPER_URLS       = "atlas.graph.index.search.solr.zookeeper-urls";
    public static final String INDEX_BACKEND_LUCENE      = "lucene";
    public static final String INDEX_BACKEND_ES          = "elasticsearch";
    public static final String GRAPH_TX_LOG_CONF         = "tx.log-tx";
    public static final String GRAPH_TX_LOG_VERBOSE_CONF = "tx.recovery.verbose";
    public static final String GRAPH_TX_LOG_TTL_CONF     = "log.tx.ttl";

    private static final String OLDER_STORAGE_EXCEPTION = "Storage version is incompatible with current client";

    private static volatile AtlasJanusGraph atlasGraphInstance;
    private static volatile JanusGraph      graphInstance;

    public AtlasJanusGraphDatabase() {
        //update registry
        GraphSONMapper.build().addRegistry(JanusGraphIoRegistry.instance()).create();
    }

    public static Configuration getConfiguration() throws AtlasException {
        Configuration configProperties = ApplicationProperties.get();

        if (isEmbeddedSolr()) { // AtlasJanusGraphIndexClient.performRequestHandlerAction() fails for embedded-solr; disable freetext until this issue is resolved
            startEmbeddedSolr();

            configProperties.setProperty(ApplicationProperties.ENABLE_FREETEXT_SEARCH_CONF, false);
        }

        configProperties.setProperty(SOLR_ZOOKEEPER_URLS, configProperties.getStringArray(SOLR_ZOOKEEPER_URL));

        Configuration janusConfig = ApplicationProperties.getSubsetConfiguration(configProperties, GRAPH_PREFIX);

        //add serializers for non-standard property value types that Atlas uses
        janusConfig.setProperty("attributes.custom.attribute1.attribute-class", TypeCategory.class.getName());
        janusConfig.setProperty("attributes.custom.attribute1.serializer-class", TypeCategorySerializer.class.getName());

        //not ideal, but avoids making large changes to Atlas
        janusConfig.setProperty("attributes.custom.attribute2.attribute-class", ArrayList.class.getName());
        janusConfig.setProperty("attributes.custom.attribute2.serializer-class", SerializableSerializer.class.getName());

        janusConfig.setProperty("attributes.custom.attribute3.attribute-class", BigInteger.class.getName());
        janusConfig.setProperty("attributes.custom.attribute3.serializer-class", BigIntegerSerializer.class.getName());

        janusConfig.setProperty("attributes.custom.attribute4.attribute-class", BigDecimal.class.getName());
        janusConfig.setProperty("attributes.custom.attribute4.serializer-class", BigDecimalSerializer.class.getName());

        return janusConfig;
    }

    public static JanusGraph getGraphInstance() {
        JanusGraph me = graphInstance;

        if (me == null) {
            synchronized (AtlasJanusGraphDatabase.class) {
                me = graphInstance;

                if (me == null) {
                    Configuration config;

                    try {
                        config = getConfiguration();
                    } catch (AtlasException e) {
                        throw new RuntimeException(e);
                    }

                    configureTxLogBasedIndexRecovery();

                    graphInstance      = initJanusGraph(config);
                    atlasGraphInstance = new AtlasJanusGraph();

                    me = graphInstance;

                    validateIndexBackend(config);
                }
            }
        }

        return me;
    }

    public static void configureTxLogBasedIndexRecovery() {
        try {
            boolean  recoveryEnabled = ApplicationProperties.get().getBoolean(INDEX_RECOVERY_CONF, DEFAULT_INDEX_RECOVERY);
            long     ttl             = AtlasConfiguration.SOLR_INDEX_TX_LOG_TTL_CONF.getLong();
            Duration txLogTtlSecs    = Duration.ofSeconds(Duration.ofHours(ttl).getSeconds());

            Map<String, Object> properties = new HashMap<>();

            properties.put(GRAPH_TX_LOG_CONF, recoveryEnabled);
            properties.put(GRAPH_TX_LOG_VERBOSE_CONF, recoveryEnabled);
            properties.put(GRAPH_TX_LOG_TTL_CONF, txLogTtlSecs);

            updateGlobalConfiguration(properties);

            LOG.info("Tx Log-based Index Recovery: {}!", recoveryEnabled ? "Enabled" : "Disabled");
        } catch (Exception e) {
            LOG.error("Error: Failed!", e);
        }
    }

    public static JanusGraph getBulkLoadingGraphInstance() {
        try {
            Configuration cfg = getConfiguration();

            cfg.setProperty("storage.batch-loading", true);

            org.apache.commons.configuration2.Configuration conf2 = createConfiguration2(cfg);

            return JanusGraphFactory.open(conf2);
        } catch (IllegalArgumentException | AtlasException ex) {
            LOG.error("getBulkLoadingGraphInstance: Failed!", ex);
        }

        return null;
    }

    public static void unload() {
        synchronized (AtlasJanusGraphDatabase.class) {
            JanusGraph me = graphInstance;

            if (me != null) {
                me.tx().commit();
                me.close();

                graphInstance = null;
            }
        }
    }

    public static boolean isEmbeddedSolr() {
        boolean ret = false;

        try {
            Configuration conf     = ApplicationProperties.get();
            Object        property = conf.getProperty("atlas.graph.index.search.solr.embedded");

            if (property instanceof String) {
                ret = Boolean.parseBoolean((String) property);
            }
        } catch (AtlasException ignored) {
            // ignore
        }

        return ret;
    }

    @VisibleForTesting
    static JanusGraph initJanusGraph(Configuration config) {
        org.apache.commons.configuration2.Configuration conf2 = createConfiguration2(config);

        try {
            return JanusGraphFactory.open(conf2);
        } catch (JanusGraphException e) {
            LOG.warn("JanusGraphException: {}", e.getMessage());

            if (e.getMessage().startsWith(OLDER_STORAGE_EXCEPTION)) {
                LOG.info("Newer client is being used with older janus storage version. Setting allow-upgrade=true and reattempting connection");

                config.addProperty("graph.allow-upgrade", true);

                return JanusGraphFactory.open(conf2);
            } else {
                throw new RuntimeException(e);
            }
        }
    }

    static void validateIndexBackend(Configuration config) {
        JanusGraphManagement managementSystem = null;

        try {
            managementSystem = getGraphInstance().openManagement();

            String configuredIndexBackend = config.getString(INDEX_BACKEND_CONF);
            String currentIndexBackend    = managementSystem.get(INDEX_BACKEND_CONF);

            if (!configuredIndexBackend.equals(currentIndexBackend)) {
                throw new RuntimeException("Configured Index Backend " + configuredIndexBackend + " differs from earlier configured Index Backend " + currentIndexBackend + ". Aborting!");
            }
        } finally {
            if (managementSystem != null) {
                managementSystem.commit();
            }
        }
    }

    @Override
    public boolean isGraphLoaded() {
        return graphInstance != null;
    }

    @Override
    public AtlasGraph<AtlasJanusVertex, AtlasJanusEdge> getGraph() {
        getGraphInstance();

        return atlasGraphInstance;
    }

    @Override
    public AtlasGraph<AtlasJanusVertex, AtlasJanusEdge> getGraphBulkLoading() {
        return new AtlasJanusGraph(getBulkLoadingGraphInstance());
    }

    @Override
    public void initializeTestGraph() {
        //nothing to do
    }

    @Override
    public void cleanup() {
        JanusGraph g = getGraphInstance();

        try {
            if (g != null) {
                g.close();
            }
        } catch (Throwable t) {
            LOG.warn("Could not close test JanusGraph", t);

            t.printStackTrace();
        }

        try {
            if (g != null) {
                JanusGraphFactory.drop(g);
            }
        } catch (Throwable t) {
            LOG.warn("Could not clear test JanusGraph", t);

            t.printStackTrace();
        }

        if (isEmbeddedSolr()) {
            try {
                stopEmbeddedSolr();
            } catch (Throwable t) {
                LOG.warn("Could not stop local solr server", t);
            }
        }
    }

    private static void addHBase2Support() {
        try {
            Field field = StandardStoreManager.class.getDeclaredField("ALL_MANAGER_CLASSES");

            field.setAccessible(true);

            Field modifiersField = Field.class.getDeclaredField("modifiers");

            modifiersField.setAccessible(true);
            modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL);

            Map<String, String> customMap = new HashMap<>(StandardStoreManager.getAllManagerClasses());

            customMap.put("hbase2", HBaseStoreManager.class.getName());

            ImmutableMap<String, String> immap = ImmutableMap.copyOf(customMap);

            field.set(null, immap);

            LOG.debug("Injected HBase2 support - {}", HBaseStoreManager.class.getName());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static void addRdbmsSupport() {
        try {
            Field field = StandardStoreManager.class.getDeclaredField("ALL_MANAGER_CLASSES");

            field.setAccessible(true);

            Field modifiersField = Field.class.getDeclaredField("modifiers");

            modifiersField.setAccessible(true);
            modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL);

            Map<String, String> customMap = new HashMap<>(StandardStoreManager.getAllManagerClasses());

            customMap.put("rdbms", RdbmsStoreManager.class.getName());

            ImmutableMap<String, String> immap = ImmutableMap.copyOf(customMap);

            field.set(null, immap);

            LOG.debug("Injected RDBMS support - {}", RdbmsStoreManager.class.getName());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static void addSolr6Index() {
        try {
            Field field = StandardIndexProvider.class.getDeclaredField("ALL_MANAGER_CLASSES");

            field.setAccessible(true);

            Field modifiersField = Field.class.getDeclaredField("modifiers");

            modifiersField.setAccessible(true);
            modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL);

            Map<String, String> customMap = new HashMap<>(StandardIndexProvider.getAllProviderClasses());

            customMap.put("solr", Solr6Index.class.getName());

            ImmutableMap<String, String> immap = ImmutableMap.copyOf(customMap);

            field.set(null, immap);

            LOG.debug("Injected solr6 index - {}", Solr6Index.class.getName());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static void addElasticSearch7Index() {
        try {
            Field field = StandardIndexProvider.class.getDeclaredField("ALL_MANAGER_CLASSES");

            field.setAccessible(true);

            Field modifiersField = Field.class.getDeclaredField("modifiers");

            modifiersField.setAccessible(true);
            modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL);

            Map<String, String> customMap = new HashMap<>(StandardIndexProvider.getAllProviderClasses());

            customMap.put("elasticsearch", ElasticSearch7Index.class.getName());

            ImmutableMap<String, String> immap = ImmutableMap.copyOf(customMap);

            field.set(null, immap);

            LOG.debug("Injected es7 index - {}", ElasticSearch7Index.class.getName());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    private static void updateGlobalConfiguration(Map<String, Object> map) {
        JanusGraph           graph            = null;
        JanusGraphManagement managementSystem = null;

        try {
            graph            = initJanusGraph(getConfiguration());
            managementSystem = graph.openManagement();

            for (Map.Entry<String, Object> entry : map.entrySet()) {
                managementSystem.set(entry.getKey(), entry.getValue());
            }

            LOG.info("Global properties updated!: {}", map);
        } catch (Exception ex) {
            LOG.error("Error updating global configuration: {}", map, ex);
        } finally {
            if (managementSystem != null) {
                managementSystem.commit();
            }

            if (graph != null) {
                graph.close();
            }
        }
    }

    private static org.apache.commons.configuration2.Configuration createConfiguration2(Configuration conf) {
        Properties properties = ConfigurationConverter.getProperties(conf);

        return org.apache.commons.configuration2.ConfigurationConverter.getConfiguration(properties);
    }

    private static void startEmbeddedSolr() throws AtlasException {
        LOG.info("==> startEmbeddedSolr()");

        try {
            Class<?> localSolrRunnerClz = Class.forName("org.apache.atlas.runner.LocalSolrRunner");
            Method   startMethod        = localSolrRunnerClz.getMethod("start");

            startMethod.invoke(null);
        } catch (Exception excp) {
            LOG.error("startEmbeddedSolr(): failed", excp);

            throw new AtlasException("startEmbeddedSolr(): failed", excp);
        }

        LOG.info("<== startEmbeddedSolr()");
    }

    private static void stopEmbeddedSolr() throws AtlasException {
        LOG.info("==> stopEmbeddedSolr()");

        try {
            Class<?> localSolrRunnerClz = Class.forName("org.apache.atlas.runner.LocalSolrRunner");
            Method   stopMethod         = localSolrRunnerClz.getMethod("stop");

            stopMethod.invoke(null);
        } catch (Exception excp) {
            LOG.error("stopEmbeddedSolr(): failed", excp);

            throw new AtlasException("stopEmbeddedSolr(): failed", excp);
        }

        LOG.info("<== stopEmbeddedSolr()");
    }

    static {
        addHBase2Support();
        addRdbmsSupport();

        addSolr6Index();

        addElasticSearch7Index();
    }
}
