/**
 * Copyright © 2016-2018 The Thingsboard Authors
 *
 * Licensed 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.thingsboard.server.service.script;

import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import delight.nashornsandbox.NashornSandbox;
import delight.nashornsandbox.NashornSandboxes;
import jdk.nashorn.api.scripting.NashornScriptEngineFactory;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.thingsboard.server.common.data.id.EntityId;

import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import javax.script.Invocable;
import javax.script.ScriptEngine;
import javax.script.ScriptException;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public abstract class AbstractNashornJsSandboxService implements JsSandboxService {

    private NashornSandbox sandbox;
    private ScriptEngine engine;
    private ExecutorService monitorExecutorService;
    private ListeningExecutorService evalExecutorService;

    private final Map<UUID, String> functionsMap = new ConcurrentHashMap<>();
    private final Map<BlackListKey, AtomicInteger> blackListedFunctions = new ConcurrentHashMap<>();

    private final Map<String, ScriptInfo> scriptKeyToInfo = new ConcurrentHashMap<>();
    private final Map<UUID, ScriptInfo> scriptIdToInfo = new ConcurrentHashMap<>();

    @PostConstruct
    public void init() {
        if (useJsSandbox()) {
            sandbox = NashornSandboxes.create();
            monitorExecutorService = Executors.newFixedThreadPool(getMonitorThreadPoolSize());
            evalExecutorService = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(10));
            sandbox.setExecutor(monitorExecutorService);
            sandbox.setMaxCPUTime(getMaxCpuTime());
            sandbox.allowNoBraces(false);
            sandbox.setMaxPreparedStatements(30);
        } else {
            NashornScriptEngineFactory factory = new NashornScriptEngineFactory();
            engine = factory.getScriptEngine(new String[]{"--no-java"});
        }
    }

    @PreDestroy
    public void stop() {
        if (monitorExecutorService != null) {
            monitorExecutorService.shutdownNow();
        }
        if (evalExecutorService != null) {
            evalExecutorService.shutdownNow();
        }
    }

    protected abstract boolean useJsSandbox();

    protected abstract int getMonitorThreadPoolSize();

    protected abstract long getMaxCpuTime();

    protected abstract int getMaxErrors();

    @Override
    public ListenableFuture<UUID> eval(JsScriptType scriptType, String scriptBody, String... argNames) {
        ScriptInfo scriptInfo = deduplicate(scriptType, scriptBody);
        UUID scriptId = scriptInfo.getId();
        AtomicInteger duplicateCount = scriptInfo.getCount();

        synchronized (scriptInfo.getLock()) {
            if (duplicateCount.compareAndSet(0, 1)) {
                try {
                    evaluate(scriptId, scriptType, scriptBody, argNames);
                } catch (Exception e) {
                    duplicateCount.decrementAndGet();
                    log.warn("Failed to compile JS script: {}", e.getMessage(), e);
                    return Futures.immediateFailedFuture(e);
                }
            } else {
                duplicateCount.incrementAndGet();
            }
        }
        return Futures.immediateFuture(scriptId);
    }

    private void evaluate(UUID scriptId, JsScriptType scriptType, String scriptBody, String... argNames) throws ScriptException {
        String functionName = "invokeInternal_" + scriptId.toString().replace('-', '_');
        String jsScript = generateJsScript(scriptType, functionName, scriptBody, argNames);
        if (useJsSandbox()) {
            sandbox.eval(jsScript);
        } else {
            engine.eval(jsScript);
        }
        functionsMap.put(scriptId, functionName);
    }

    @Override
    public ListenableFuture<Object> invokeFunction(UUID scriptId, EntityId entityId, Object... args) {
        String functionName = functionsMap.get(scriptId);
        if (functionName == null) {
            return Futures.immediateFailedFuture(new RuntimeException("No compiled script found for scriptId: [" + scriptId + "]!"));
        }
        if (!isBlackListed(scriptId)) {
            try {
                Object result;
                if (useJsSandbox()) {
                    result = sandbox.getSandboxedInvocable().invokeFunction(functionName, args);
                } else {
                    result = ((Invocable) engine).invokeFunction(functionName, args);
                }
                return Futures.immediateFuture(result);
            } catch (Exception e) {
                BlackListKey blackListKey = new BlackListKey(scriptId, entityId);
                blackListedFunctions.computeIfAbsent(blackListKey, key -> new AtomicInteger(0)).incrementAndGet();
                return Futures.immediateFailedFuture(e);
            }
        } else {
            return Futures.immediateFailedFuture(
                    new RuntimeException("Script is blacklisted due to maximum error count " + getMaxErrors() + "!"));
        }
    }

    @Override
    public ListenableFuture<Void> release(UUID scriptId, EntityId entityId) {
        ScriptInfo scriptInfo = scriptIdToInfo.get(scriptId);
        if (scriptInfo == null) {
            log.warn("Script release called for not existing script id [{}]", scriptId);
            return Futures.immediateFuture(null);
        }

        synchronized (scriptInfo.getLock()) {
            int remainingDuplicates = scriptInfo.getCount().decrementAndGet();
            if (remainingDuplicates > 0) {
                return Futures.immediateFuture(null);
            }

            String functionName = functionsMap.get(scriptId);
            if (functionName != null) {
                try {
                    if (useJsSandbox()) {
                        sandbox.eval(functionName + " = undefined;");
                    } else {
                        engine.eval(functionName + " = undefined;");
                    }
                    functionsMap.remove(scriptId);
                    blackListedFunctions.remove(new BlackListKey(scriptId, entityId));
                } catch (ScriptException e) {
                    log.error("Could not release script [{}] [{}]", scriptId, remainingDuplicates);
                    return Futures.immediateFailedFuture(e);
                }
            } else {
                log.warn("Function name do not exist for script [{}] [{}]", scriptId, remainingDuplicates);
            }
        }
        return Futures.immediateFuture(null);
    }


    private boolean isBlackListed(UUID scriptId) {
        if (blackListedFunctions.containsKey(scriptId)) {
            AtomicInteger errorCount = blackListedFunctions.get(scriptId);
            return errorCount.get() >= getMaxErrors();
        } else {
            return false;
        }
    }

    private String generateJsScript(JsScriptType scriptType, String functionName, String scriptBody, String... argNames) {
        switch (scriptType) {
            case RULE_NODE_SCRIPT:
                return RuleNodeScriptFactory.generateRuleNodeScript(functionName, scriptBody, argNames);
            default:
                throw new RuntimeException("No script factory implemented for scriptType: " + scriptType);
        }
    }

    private ScriptInfo deduplicate(JsScriptType scriptType, String scriptBody) {
        ScriptInfo meta = ScriptInfo.preInit();
        String key = deduplicateKey(scriptType, scriptBody);
        ScriptInfo latestMeta = scriptKeyToInfo.computeIfAbsent(key, i -> meta);
        return scriptIdToInfo.computeIfAbsent(latestMeta.getId(), i -> latestMeta);
    }

    private String deduplicateKey(JsScriptType scriptType, String scriptBody) {
        return scriptType + "_" + scriptBody;
    }

    @Getter
    private static class ScriptInfo {
        private final UUID id;
        private final Object lock;
        private final AtomicInteger count;

        ScriptInfo(UUID id, Object lock, AtomicInteger count) {
            this.id = id;
            this.lock = lock;
            this.count = count;
        }

        static ScriptInfo preInit() {
            UUID preId = UUID.randomUUID();
            AtomicInteger preCount = new AtomicInteger();
            Object preLock = new Object();
            return new ScriptInfo(preId, preLock, preCount);
        }
    }

    @EqualsAndHashCode
    @Getter
    private static class BlackListKey {
        private final UUID scriptId;
        private final EntityId entityId;

        public BlackListKey(UUID scriptId, EntityId entityId) {
            this.scriptId = scriptId;
            this.entityId = entityId;
        }
    }
}
