/*
 * Copyright 2017 LinkedIn Corp.
 *
 * 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 azkaban.execapp;

import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import azkaban.event.Event;
import azkaban.execapp.event.FlowWatcher;
import azkaban.execapp.jmx.JmxJobMBeanManager;
import azkaban.executor.ExecutableFlow;
import azkaban.executor.ExecutionOptions;
import azkaban.executor.ExecutionOptions.FailureAction;
import azkaban.executor.ExecutorLoader;
import azkaban.executor.InteractiveTestJob;
import azkaban.executor.MockExecutorLoader;
import azkaban.executor.Status;
import azkaban.flow.Flow;
import azkaban.jobtype.JobTypeManager;
import azkaban.jobtype.JobTypePluginSet;
import azkaban.project.FlowLoader;
import azkaban.project.FlowLoaderFactory;
import azkaban.project.Project;
import azkaban.project.ProjectLoader;
import azkaban.project.ProjectManagerException;
import azkaban.test.Utils;
import azkaban.test.executions.ExecutionsTestUtil;
import azkaban.utils.JSONUtils;
import azkaban.utils.Props;
import java.io.File;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import org.apache.commons.io.FileUtils;
import org.junit.rules.TemporaryFolder;

public class FlowRunnerTestUtil {

  private static int id = 101;
  private final Map<String, Flow> flowMap;
  private final Project project;
  private final File workingDir;
  private final JobTypeManager jobtypeManager;
  private final File projectDir;
  private final ProjectLoader projectLoader;
  private ExecutorLoader executorLoader;

  public FlowRunnerTestUtil(final String flowName, final TemporaryFolder temporaryFolder)
      throws Exception {

    this.projectDir = ExecutionsTestUtil.getFlowDir(flowName);
    this.workingDir = temporaryFolder.newFolder();
    this.project = new Project(1, "testProject");

    this.flowMap = FlowRunnerTestUtil
        .prepareProject(this.project, this.projectDir, this.workingDir);

    this.executorLoader = mock(ExecutorLoader.class);
    this.projectLoader = mock(ProjectLoader.class);
    when(this.executorLoader.updateExecutableReference(anyInt(), anyLong())).thenReturn(true);

    Utils.initServiceProvider();
    JmxJobMBeanManager.getInstance().initialize(new Props());

    InteractiveTestJob.clearTestJobs();

    this.jobtypeManager = new JobTypeManager(null, null, this.getClass().getClassLoader());
    final JobTypePluginSet pluginSet = this.jobtypeManager.getJobTypePluginSet();
    pluginSet.addPluginClass("test", InteractiveTestJob.class);
  }

  /**
   * Initialize the project with the flow definitions stored in the given source directory. Also
   * copy the source directory to the working directory.
   *
   * @param project project to initialize
   * @param sourceDir the source dir
   * @param workingDir the working dir
   * @return the flow name to flow map
   * @throws ProjectManagerException the project manager exception
   * @throws IOException the io exception
   */
  public static Map<String, Flow> prepareProject(final Project project, final File sourceDir,
      final File workingDir)
      throws ProjectManagerException, IOException {
    final FlowLoaderFactory loaderFactory = new FlowLoaderFactory(new Props(null));
    final FlowLoader loader = loaderFactory.createFlowLoader(sourceDir);

    loader.loadProjectFlow(project, sourceDir);
    if (!loader.getErrors().isEmpty()) {
      for (final String error : loader.getErrors()) {
        System.out.println(error);
      }
      throw new RuntimeException(String.format(
          "Errors found in loading flows into a project ( %s ). From the directory: ( %s ).",
          project.getName(), sourceDir));
    }

    final Map<String, Flow> flowMap = loader.getFlowMap();
    project.setFlows(flowMap);
    FileUtils.copyDirectory(sourceDir, workingDir);
    return flowMap;
  }

  public static void waitEventFired(final EventCollectorListener eventCollector,
      final String nestedId, final Status status)
      throws InterruptedException {
    for (int i = 0; i < 1000; i++) {
      for (final Event event : eventCollector.getEventList()) {
        if (event.getData().getStatus() == status && event.getData().getNestedId()
            .equals(nestedId)) {
          return;
        }
      }
      synchronized (EventCollectorListener.handleEvent) {
        EventCollectorListener.handleEvent.wait(10L);
      }
    }
    fail("Event wasn't fired with [" + nestedId + "], " + status);
  }

  public static ExecutableFlow prepareExecDir(final File workingDir, final File execDir,
      final String flowName, final int execId) throws IOException {
    FileUtils.copyDirectory(execDir, workingDir);
    final File jsonFlowFile = new File(workingDir, flowName + ".flow");
    final Object flowObj = JSONUtils.parseJSONFromFile(jsonFlowFile);
    final Project project = new Project(1, "test");
    final Flow flow = Flow.flowFromObject(flowObj);
    final ExecutableFlow execFlow = new ExecutableFlow(project, flow);
    execFlow.setExecutionId(execId);
    execFlow.setExecutionPath(workingDir.getPath());
    return execFlow;
  }

  public static void startThread(final FlowRunner runner) {
    new Thread(runner).start();
  }

  public FlowRunner createFromFlowFile(final String flowName) throws Exception {
    return createFromFlowFile(new EventCollectorListener(), flowName);
  }

  public FlowRunner createFromFlowFile(final EventCollectorListener eventCollector,
      final String flowName) throws Exception {
    return createFromFlowFile(flowName, eventCollector, new ExecutionOptions());
  }

  public FlowRunner createFromFlowFile(final String flowName,
      final EventCollectorListener eventCollector, final ExecutionOptions options)
      throws Exception {
    return createFromFlowFile(flowName, eventCollector, options, null, null);
  }

  public FlowRunner createFromFlowFile(final String flowName, final FlowWatcher watcher,
      final Integer pipeline) throws Exception {
    return createFromFlowFile(flowName, new EventCollectorListener(), new ExecutionOptions(),
        watcher, pipeline);
  }

  public FlowRunner createFromFlowFile(final String flowName,
      final EventCollectorListener eventCollector,
      final ExecutionOptions options, final FlowWatcher watcher, final Integer pipeline)
      throws Exception {
    final ExecutableFlow exFlow = FlowRunnerTestUtil
        .prepareExecDir(this.workingDir, this.projectDir, flowName, 1);
    if (watcher != null) {
      options.setPipelineLevel(pipeline);
      options.setPipelineExecutionId(watcher.getExecId());
    }
    final FlowRunner runner = createFromExecutableFlow(eventCollector, exFlow, options,
        new HashMap<>(),
        new Props());
    runner.setFlowWatcher(watcher);
    return runner;
  }

  public FlowRunner createFromFlowMap(final String flowName, final String jobIdPrefix)
      throws Exception {
    return createFromFlowMap(flowName, jobIdPrefix, new ExecutionOptions());
  }

  public FlowRunner createFromFlowMap(final String flowName, final String jobIdPrefix,
      final ExecutionOptions options)
      throws Exception {
    return createFromFlowMap(flowName, jobIdPrefix, options, new Props());
  }

  public FlowRunner createFromFlowMap(final String flowName,
      final HashMap<String, String> flowParams) throws Exception {
    return createFromFlowMap(null, flowName, null, flowParams, new Props());
  }

  public FlowRunner createFromFlowMap(final String flowName, final ExecutionOptions options,
      final Map<String, String> flowParams, final Props azkabanProps) throws Exception {
    return createFromFlowMap(null, flowName, options, flowParams,
        azkabanProps);
  }

  public FlowRunner createFromFlowMap(final EventCollectorListener eventCollector,
      final String flowName, final ExecutionOptions options,
      final Map<String, String> flowParams, final Props azkabanProps)
      throws Exception {
    final Flow flow = this.flowMap.get(flowName);
    final ExecutableFlow exFlow = new ExecutableFlow(this.project, flow);
    return createFromExecutableFlow(eventCollector, exFlow, options, flowParams,
        azkabanProps);
  }

  public FlowRunner createFromFlowMap(final String flowName, final FailureAction action)
      throws Exception {
    final ExecutionOptions options = new ExecutionOptions();
    options.setFailureAction(action);
    return createFromFlowMap(flowName, options, new HashMap<>(), new Props());
  }

  private FlowRunner createFromFlowMap(final String flowName, final String jobIdPrefix,
      final ExecutionOptions options, final Props azkabanProps) throws Exception {
    final Map<String, String> flowParams = new HashMap<>();
    flowParams.put(InteractiveTestJob.JOB_ID_PREFIX, jobIdPrefix);
    return createFromFlowMap(new EventCollectorListener(), flowName, options, flowParams,
        azkabanProps);
  }

  private FlowRunner createFromExecutableFlow(final EventCollectorListener eventCollector,
      final ExecutableFlow exFlow, final ExecutionOptions options,
      final Map<String, String> flowParams, final Props azkabanProps)
      throws Exception {
    final int exId = id++;
    exFlow.setExecutionPath(this.workingDir.getPath());
    exFlow.setExecutionId(exId);
    if (options != null) {
      exFlow.setExecutionOptions(options);
    }
    exFlow.getExecutionOptions().addAllFlowParameters(flowParams);
    this.executorLoader.uploadExecutableFlow(exFlow);
    final FlowRunner runner =
        new FlowRunner(exFlow, this.executorLoader, this.projectLoader,
            this.jobtypeManager, azkabanProps, null);
    if (eventCollector != null) {
      runner.addListener(eventCollector);
    }
    return runner;
  }

  public ExecutorLoader getExecutorLoader() {
    return this.executorLoader;
  }

  public void setExecutorLoader(final MockExecutorLoader executorLoader) {
    this.executorLoader = executorLoader;
  }

  public ProjectLoader getProjectLoader() {
    return this.projectLoader;
  }

  public Project getProject() {
    return this.project;
  }
}
