/*
 * 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
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.tinkerpop.gremlin.server.util;

import com.codahale.metrics.Gauge;
import io.netty.channel.Channel;
import io.netty.channel.group.DefaultChannelGroup;
import io.netty.util.concurrent.GlobalEventExecutor;
import org.apache.tinkerpop.gremlin.groovy.engine.GremlinExecutor;
import org.apache.tinkerpop.gremlin.jsr223.GremlinScriptEngine;
import org.apache.tinkerpop.gremlin.process.traversal.TraversalSource;
import org.apache.tinkerpop.gremlin.server.Channelizer;
import org.apache.tinkerpop.gremlin.server.GraphManager;
import org.apache.tinkerpop.gremlin.server.GremlinServer;
import org.apache.tinkerpop.gremlin.server.Settings;
import org.apache.tinkerpop.gremlin.server.transaction.TransactionManager;
import org.apache.tinkerpop.gremlin.structure.Graph;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.script.SimpleBindings;
import java.lang.reflect.Constructor;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static com.codahale.metrics.MetricRegistry.name;

/**
 * The core of script execution in Gremlin Server.  Given {@link Settings} and optionally other arguments, this
 * class will construct a {@link GremlinExecutor} to be used by Gremlin Server.  A typical usage would be to
 * instantiate the {@link GremlinServer} and then immediately call {@link GremlinServer#getServerGremlinExecutor()}
 * which would allow the opportunity to assign "host options" which could be used by a custom {@link Channelizer}.
 * Add these options before calling {@link GremlinServer#start()} to be sure the {@link Channelizer} gets access to
 * those.
 *
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */
public class ServerGremlinExecutor {
    private static final Logger logger = LoggerFactory.getLogger(ServerGremlinExecutor.class);

    private final GraphManager graphManager;
    private final TransactionManager transactionManager;
    private final Settings settings;
    private final List<LifeCycleHook> hooks;

    private final ScheduledExecutorService scheduledExecutorService;
    private final ExecutorService gremlinExecutorService;
    private final GremlinExecutor gremlinExecutor;
    private final DefaultChannelGroup channels = new DefaultChannelGroup(GlobalEventExecutor.INSTANCE);

    private final Map<String,Object> hostOptions = new ConcurrentHashMap<>();

    /**
     * The total number of channels.
     */
    public final Gauge channelsTotalGauge = MetricManager.INSTANCE.getGauge(
            this::getChannelCount, name(GremlinServer.class, "channels", "total"));

    /**
     * The total number of paused channels.
     */
    public final Gauge channelsPausedGauge = MetricManager.INSTANCE.getGauge(
            this::getPausedChannelCount, name(GremlinServer.class, "channels", "paused"));

    /**
     * Create a new object from {@link Settings} where thread pools are externally assigned. Note that if the
     * {@code scheduleExecutorServiceClass} is set to {@code null} it will be created via
     * {@link Executors#newScheduledThreadPool(int, ThreadFactory)}.  If either of the {@link ExecutorService}
     * instances are supplied, the {@link Settings#gremlinPool} value will be ignored for the pool size.
     */
    public ServerGremlinExecutor(final Settings settings, final ExecutorService gremlinExecutorService,
                                 final ScheduledExecutorService scheduledExecutorService) {
        this.settings = settings;

        try {
            final Class<?> clazz = Class.forName(settings.graphManager);
            final Constructor c = clazz.getConstructor(Settings.class);
            graphManager = (GraphManager) c.newInstance(settings);
        } catch (ClassNotFoundException e) {
            logger.error("Could not find GraphManager implementation "
                         + "defined by the 'graphManager' setting as: {}",
                         settings.graphManager);
            throw new RuntimeException(e);
        } catch (Exception e) {
            logger.error("Could not invoke constructor on class {} (defined by "
                         + "the 'graphManager' setting) with one argument of "
                         + "class Settings",
                         settings.graphManager);
            throw new RuntimeException(e);
        }

        if (null == gremlinExecutorService) {
            final ThreadFactory threadFactoryGremlin = ThreadFactoryUtil.create("exec-%d");
            final BlockingQueue<Runnable> queue = new ArrayBlockingQueue<>(settings.maxWorkQueueSize);
            this.gremlinExecutorService = new ThreadPoolExecutor(settings.gremlinPool, settings.gremlinPool,
                    0L, TimeUnit.MILLISECONDS, queue, threadFactoryGremlin,
                    new ThreadPoolExecutor.AbortPolicy());
        } else {
            this.gremlinExecutorService = gremlinExecutorService;
        }

        if (null == scheduledExecutorService) {
            final ThreadFactory threadFactoryGremlin = ThreadFactoryUtil.create("worker-%d");
            this.scheduledExecutorService = Executors.newScheduledThreadPool(settings.threadPoolWorker, threadFactoryGremlin);
        } else {
            this.scheduledExecutorService = scheduledExecutorService;
        }

        logger.info("Initialized Gremlin thread pool.  Threads in pool named with pattern gremlin-*");

        // build the allowlist of script engines from the server config - gremlin-lang is always available
        final Set<String> allowedEngines = new HashSet<>(settings.scriptEngines.keySet());
        allowedEngines.add("gremlin-lang");

        final GremlinExecutor.Builder gremlinExecutorBuilder = GremlinExecutor.build()
                .timeoutMillis(settings.getTimeoutMillis())
                .afterFailure((b, e) -> this.graphManager.rollbackAll())
                .beforeEval(b -> this.graphManager.rollbackAll())
                .afterTimeout((b, e) -> this.graphManager.rollbackAll())
                .globalBindings(this.graphManager.getAsBindings())
                .executorService(this.gremlinExecutorService)
                .scheduledExecutorService(this.scheduledExecutorService)
                .allowedEngineNames(allowedEngines);

        settings.scriptEngines.forEach((k, v) -> {
            // use plugins if they are present
            if (!v.plugins.isEmpty()) {
                // make sure that server related classes are available at init - new approach. the LifeCycleHook stuff
                // will be added explicitly via configuration using GremlinServerGremlinModule in the yaml
                gremlinExecutorBuilder.addPlugins(k, v.plugins);
            }
        });

        gremlinExecutor = gremlinExecutorBuilder.create();

        logger.info("Initialized GremlinExecutor and preparing GremlinScriptEngines instances.");

        // force each scriptengine to process something so that the init scripts will fire (this is necessary if
        // the GremlinExecutor is using the GremlinScriptEngineManager. this is a bit of hack, but it at least allows
        // the global bindings to become available after the init scripts are run (DefaultGremlinScriptEngineManager
        // runs the init scripts when the GremlinScriptEngine is created.
        settings.scriptEngines.keySet().forEach(engineName -> {
            try {
                // gremlin-lang does not need to be initialized. not so nice, but gremlin-lang is the only exception
                // and ultimately, gremlin-lang will likely end up the only choice in gremlin-server.
                // groovy-test (used for testing purposes only) also interprets input via the Gremlin grammar, so it
                // cannot evaluate the "1+1" init probe and likewise does not need init-script priming (its engine
                // setup happens at construction, which is triggered by registerMetrics below).
                if (!engineName.equals("gremlin-lang") && !engineName.equals("groovy-test")) {
                    // use no timeout on the engine initialization - perhaps this can be a configuration later
                    final GremlinExecutor.LifeCycle lifeCycle = GremlinExecutor.LifeCycle.build().
                            timeoutMillisOverride(0L).create();
                    gremlinExecutor.eval("1+1", engineName, new SimpleBindings(Collections.emptyMap()), lifeCycle).join();
                }

                registerMetrics(engineName);
                logger.info("Initialized {} GremlinScriptEngine and registered metrics", engineName);
            } catch (Exception ex) {
                logger.warn(String.format("Could not initialize %s GremlinScriptEngine as init script could not be evaluated", engineName), ex);
            }
        });

        // script engine init may have altered the graph bindings or maybe even created new ones - need to
        // re-apply those references back
        gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
                .filter(kv -> kv.getValue() instanceof Graph)
                .forEach(kv -> this.graphManager.putGraph(kv.getKey(), (Graph) kv.getValue()));

        // script engine init may have constructed the TraversalSource bindings - store them in Graphs object
        final boolean hasScriptTraversalSources = gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
                .anyMatch(kv -> kv.getValue() instanceof TraversalSource);
        if (hasScriptTraversalSources) {
            logger.warn("TraversalSource instances were created via script engine initialization scripts. " +
                        "This approach is deprecated - use the 'traversalSources' section within 'graphs' YAML configuration instead.");
            gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
                    .filter(kv -> kv.getValue() instanceof TraversalSource)
                    .forEach(kv -> {
                        logger.info("A {} is now bound to [{}] with {}", kv.getValue().getClass().getSimpleName(), kv.getKey(), kv.getValue());
                        this.graphManager.putTraversalSource(kv.getKey(), (TraversalSource) kv.getValue());
                    });
        }

        // determine if the initialization scripts introduced LifeCycleHook objects - if so we need to gather them
        // up for execution
        final boolean hasScriptHooks = gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
                .anyMatch(kv -> kv.getValue() instanceof LifeCycleHook);
        if (hasScriptHooks) {
            logger.warn("LifeCycleHook instances were created via script engine initialization scripts. " +
                        "This approach is deprecated - use the 'lifecycleHooks' YAML configuration instead.");
        }
        final List<LifeCycleHook> scriptHooks = gremlinExecutor.getScriptEngineManager().getBindings().entrySet().stream()
                .filter(kv -> kv.getValue() instanceof LifeCycleHook)
                .map(kv -> (LifeCycleHook) kv.getValue())
                .collect(Collectors.toList());

        // process declarative traversalSources from per-graph YAML config
        final Set<String> graphsWithExplicitTraversalSources = new HashSet<>();
        for (final String graphName : settings.graphs.keySet()) {
            final Settings.GraphSettings graphSettings = settings.getGraphSettings(graphName);
            if (graphSettings == null || graphSettings.traversalSources == null || graphSettings.traversalSources.isEmpty())
                continue;

            final Graph graph = this.graphManager.getGraph(graphName);
            if (null == graph) {
                logger.warn("Could not create TraversalSources for graph [{}] - graph not found", graphName);
                continue;
            }

            graphsWithExplicitTraversalSources.add(graphName);

            for (final Settings.TraversalSourceSettings tsSettings : graphSettings.traversalSources) {
                final String tsName = tsSettings.name;
                if (tsSettings.gremlinExpression != null && !tsSettings.gremlinExpression.isEmpty()) {
                    // resolve which script engine to use for the expression
                    final String language = resolveLanguage(tsSettings.language);
                    try {
                        // bind a base traversal source as 'g' for the expression to operate on
                        final SimpleBindings bindings = new SimpleBindings();
                        bindings.put("g", graph.traversal());
                        final TraversalSource ts = (TraversalSource) gremlinExecutor.eval(
                                tsSettings.gremlinExpression, language, bindings).join();
                        this.graphManager.putTraversalSource(tsName, ts);
                        logger.info("A {} is now bound to [{}] via gremlinExpression", ts.getClass().getSimpleName(), tsName);
                    } catch (Exception ex) {
                        logger.warn(String.format("Could not create TraversalSource [%s] from gremlinExpression - %s",
                                    tsName, ex.getMessage()), ex);
                    }
                } else {
                    final TraversalSource ts = graph.traversal();
                    this.graphManager.putTraversalSource(tsName, ts);
                    logger.info("A {} is now bound to [{}]", ts.getClass().getSimpleName(), tsName);
                }
            }
        }

        // auto-create TraversalSources for graphs that don't have explicit traversalSources entries and
        // were not already bound by script engine initialization
        for (final String graphName : this.graphManager.getGraphNames()) {
            if (graphsWithExplicitTraversalSources.contains(graphName))
                continue;

            final String tsName = graphName.equals("graph") ? "g" : "g_" + graphName;
            if (this.graphManager.getTraversalSource(tsName) != null)
                continue;

            final Graph graph = this.graphManager.getGraph(graphName);
            final TraversalSource ts = graph.traversal();
            this.graphManager.putTraversalSource(tsName, ts);
            logger.info("A {} is now auto-bound to [{}] for graph [{}]",
                        ts.getClass().getSimpleName(), tsName, graphName);
        }

        // instantiate Java-based lifecycle hooks from YAML config
        final List<LifeCycleHook> yamlHooks = new ArrayList<>();
        if (settings.lifecycleHooks != null) {
            for (final Settings.LifeCycleHookSettings hookSettings : settings.lifecycleHooks) {
                try {
                    final Class<?> clazz = Class.forName(hookSettings.className);
                    final LifeCycleHook hook = (LifeCycleHook) clazz.getDeclaredConstructor().newInstance();
                    if (hookSettings.config != null) {
                        hook.init(hookSettings.config);
                    }
                    yamlHooks.add(hook);
                    logger.info("Instantiated LifeCycleHook: {}", hookSettings.className);
                } catch (Exception ex) {
                    logger.error(String.format("Could not instantiate LifeCycleHook %s", hookSettings.className), ex);
                }
            }
        }

        // combine YAML-configured hooks with any script-produced hooks (YAML hooks execute first)
        final List<LifeCycleHook> allHooks = new ArrayList<>(yamlHooks);
        allHooks.addAll(scriptHooks);
        hooks = allHooks;

        transactionManager = new TransactionManager(
                scheduledExecutorService,
                graphManager,
                settings.idleTransactionTimeoutMillis,
                settings.maxTransactionLifetimeMillis,
                settings.maxConcurrentTransactions,
                settings.perGraphCloseTimeoutMillis
        );
    }

    private void registerMetrics(final String engineName) {
        final GremlinScriptEngine engine = gremlinExecutor.getScriptEngineManager().getEngineByName(engineName);
        MetricManager.INSTANCE.registerGremlinScriptEngineMetrics(engine, engineName, "sessionless", "class-cache");
    }

    /**
     * Resolves the script engine language to use for evaluating a traversal source gremlinExpression. If an explicit language
     * is provided, it is used directly. Otherwise, if exactly one script engine is configured, that engine is used.
     * Falls back to {@code gremlin-lang}.
     */
    private String resolveLanguage(final String explicitLanguage) {
        if (explicitLanguage != null && !explicitLanguage.isEmpty())
            return explicitLanguage;

        if (settings.scriptEngines.size() == 1)
            return settings.scriptEngines.keySet().iterator().next();

        return "gremlin-lang";
    }

    public void addHostOption(final String key, final Object value) {
        hostOptions.put(key, value);
    }

    public Map<String,Object> getHostOptions() {
        return Collections.unmodifiableMap(hostOptions);
    }

    public Object removeHostOption(final String key) {
        return hostOptions.remove(key);
    }

    public void clearHostOptions() {
        hostOptions.clear();
    }

    public ScheduledExecutorService getScheduledExecutorService() {
        return scheduledExecutorService;
    }

    public GremlinExecutor getGremlinExecutor() {
        return gremlinExecutor;
    }

    public ExecutorService getGremlinExecutorService() {
        return gremlinExecutorService;
    }

    public GraphManager getGraphManager() {
        return graphManager;
    }

    public TransactionManager getTransactionManager() {
        return transactionManager;
    }

    public Settings getSettings() {
        return settings;
    }

    public List<LifeCycleHook> getHooks() {
        return hooks;
    }

    /**
     * Gets all open channels (excluding the {@code ServerSocketChannel}).
     */
    public DefaultChannelGroup getChannels() {
        return channels;
    }

    /**
     * Gets the number of open channels currently in a paused state that are not writeable.
     */
    public long getPausedChannelCount() {
        return channels.stream().filter(ch -> !ch.isWritable()).count();
    }

    /**
     * Gets the number of open channels.
     */
    public long getChannelCount() {
        return channels.size();
    }

}
