/*
 * 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.zeppelin.mongodb;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.Scanner;

import org.apache.commons.exec.CommandLine;
import org.apache.commons.exec.DefaultExecutor;
import org.apache.commons.exec.ExecuteException;
import org.apache.commons.exec.ExecuteWatchdog;
import org.apache.commons.exec.Executor;
import org.apache.commons.exec.PumpStreamHandler;
import org.apache.commons.io.FileUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.zeppelin.interpreter.Interpreter;
import org.apache.zeppelin.interpreter.InterpreterContext;
import org.apache.zeppelin.interpreter.InterpreterResult;
import org.apache.zeppelin.interpreter.InterpreterResult.Code;
import org.apache.zeppelin.scheduler.Scheduler;
import org.apache.zeppelin.scheduler.SchedulerFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * MongoDB interpreter. It uses the mongo shell to interpret the commands.
 *
 */
public class MongoDbInterpreter extends Interpreter {

  private static final Logger LOGGER = LoggerFactory.getLogger(MongoDbInterpreter.class);

  private static final String SHELL_EXTENSION =
    new Scanner(MongoDbInterpreter.class.getResourceAsStream("/shell_extension.js"), "UTF-8")
      .useDelimiter("\\A").next();

  private long commandTimeout = 60000;

  private String dbAddress;

  private Map<String, Executor> runningProcesses = new HashMap<>();


  public MongoDbInterpreter(Properties property) {
    super(property);
  }

  @Override
  public void open() {
    commandTimeout = Long.parseLong(getProperty("mongo.shell.command.timeout"));
    dbAddress = getProperty("mongo.server.host") + ":"
      + getProperty("mongo.server.port") + "/"
      + getProperty("mongo.server.database");
  }

  @Override
  public void close() {
    //Nothing to do
  }

  @Override
  public FormType getFormType() {
    return FormType.SIMPLE;
  }

  @Override
  public InterpreterResult interpret(String script, InterpreterContext context) {
    LOGGER.debug("Run MongoDB script: {}", script);

    if (StringUtils.isEmpty(script)) {
      return new InterpreterResult(Code.SUCCESS);
    }

    // Write script in a temporary file
    // The script is enriched with extensions
    final File scriptFile = new File(getScriptFileName(context.getParagraphId()));
    try {
      FileUtils.write(scriptFile,
        SHELL_EXTENSION
          .replace("__ZEPPELIN_TABLE_LIMIT__", getProperty("mongo.shell.command.table.limit")) +
          script);
    }
    catch (IOException e) {
      LOGGER.error("Can not write script in temp file", e);
      return new InterpreterResult(Code.ERROR, e.getMessage());
    }

    InterpreterResult result = new InterpreterResult(InterpreterResult.Code.SUCCESS);;

    final DefaultExecutor executor = new DefaultExecutor();
    final ByteArrayOutputStream errorStream = new ByteArrayOutputStream();
    executor.setStreamHandler(new PumpStreamHandler(context.out, errorStream));
    executor.setWatchdog(new ExecuteWatchdog(commandTimeout)); //ExecuteWatchdog.INFINITE_TIMEOUT

    final CommandLine cmdLine = CommandLine.parse(getProperty("mongo.shell.path"));
    cmdLine.addArgument("--quiet", false);

    if (!StringUtils.isEmpty(getProperty("mongo.server.username"))) {
      cmdLine.addArgument("-u", false);
      cmdLine.addArgument(getProperty("mongo.server.username"), false);
      cmdLine.addArgument("-p", false);
      cmdLine.addArgument(getProperty("mongo.server.password"), false);
      
      if (!StringUtils.isEmpty(getProperty("mongo.server.authentdatabase"))) {
        cmdLine.addArgument("--authenticationDatabase", false);
        cmdLine.addArgument(getProperty("mongo.server.authentdatabase"), false);
      }
    }

    cmdLine.addArgument(dbAddress, false);
    cmdLine.addArgument(scriptFile.getAbsolutePath(), false);

    try {
      executor.execute(cmdLine);
      runningProcesses.put(context.getParagraphId(), executor);
    }
    catch (ExecuteException e) {
      LOGGER.error("Can not run script in paragraph {}", context.getParagraphId(), e);

      final int exitValue = e.getExitValue();
      Code code = Code.ERROR;
      String msg = errorStream.toString();
      if (exitValue == 143) {
        code = Code.INCOMPLETE;
        msg = msg + "Paragraph received a SIGTERM.\n";
        LOGGER.info("The paragraph {} stopped executing: {}", context.getParagraphId(), msg);
      }
      msg += "ExitValue: " + exitValue;
      result = new InterpreterResult(code, msg);
    }
    catch (IOException e) {
      LOGGER.error("Can not run script in paragraph {}", context.getParagraphId(), e);
      result = new InterpreterResult(Code.ERROR, e.getMessage());
    }
    finally {
      FileUtils.deleteQuietly(scriptFile);
    }

    return result;
  }

  @Override
  public int getProgress(InterpreterContext context) {
    return 0;
  }

  @Override
  public void cancel(InterpreterContext context) {
    stopProcess(context.getParagraphId());
    FileUtils.deleteQuietly(new File(getScriptFileName(context.getParagraphId())));
  }

  @Override
  public Scheduler getScheduler() {
    return SchedulerFactory.singleton().createOrGetParallelScheduler("mongo", 10);
  }

  private String getScriptFileName(String paragraphId) {
    return System.getProperty("java.io.tmpdir") + File.separator +
        "zeppelin-mongo-" + paragraphId + ".js";
  }

  private void stopProcess(String paragraphId) {
    if (runningProcesses.containsKey(paragraphId)) {
      final Executor executor = runningProcesses.get(paragraphId);
      final ExecuteWatchdog watchdog = executor.getWatchdog();
      watchdog.destroyProcess();
    }
  }
  
}