/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.master.job.workflow;

import alluxio.collections.ConcurrentHashSet;
import alluxio.exception.JobDoesNotExistException;
import alluxio.exception.status.ResourceExhaustedException;
import alluxio.job.ErrorUtils;
import alluxio.job.JobConfig;
import alluxio.job.plan.meta.PlanInfo;
import alluxio.job.wire.JobInfo;
import alluxio.job.wire.Status;
import alluxio.job.wire.WorkflowInfo;
import alluxio.job.workflow.WorkflowConfig;
import alluxio.job.workflow.WorkflowExecution;
import alluxio.job.workflow.WorkflowExecutionRegistry;
import alluxio.master.job.JobMaster;

import org.apache.commons.compress.utils.Lists;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * The {@link WorkflowTracker}.
 */
public class WorkflowTracker {
  private static final Logger LOG = LoggerFactory.getLogger(WorkflowTracker.class);

  private final JobMaster mJobMaster;

  private final ConcurrentHashMap<Long, WorkflowExecution> mWorkflows;
  private final ConcurrentHashMap<Long, ConcurrentHashSet<Long>> mWaitingOn;
  private final ConcurrentHashMap<Long, ConcurrentHashSet<Long>> mChildren;
  private final ConcurrentHashMap<Long, Long> mParentWorkflow;

  /**
   * Constructor for {@link WorkflowTracker}.
   * @param jobMaster the job master
   */
  public WorkflowTracker(JobMaster jobMaster) {
    mJobMaster = jobMaster;
    mWorkflows = new ConcurrentHashMap<>();
    mWaitingOn = new ConcurrentHashMap<>();
    mChildren = new ConcurrentHashMap<>();
    mParentWorkflow = new ConcurrentHashMap<>();
  }

  /**
   * Runs a workflow with the given configuration and job id.
   *
   * @param workflowConfig the workflow configuration
   * @param jobId the job id
   */
  public synchronized void run(WorkflowConfig workflowConfig, long jobId)
      throws JobDoesNotExistException, ResourceExhaustedException {
    WorkflowExecution workflowExecution =
        WorkflowExecutionRegistry.INSTANCE.getExecution(workflowConfig);

    mWorkflows.put(jobId, workflowExecution);

    next(jobId);
  }

  /**
   * Cancels a job with a particular job id.
   * @param jobId the job id
   * @return true if job exists, false otherwise
   */
  public synchronized boolean cancel(long jobId) {
    ConcurrentHashSet<Long> children = mChildren.get(jobId);
    if (children == null) {
      return false;
    }

    for (Long child : children) {
      try {
        mJobMaster.cancel(child);
      } catch (JobDoesNotExistException e) {
        LOG.info("Tried to cancel jobId: {} but the job did not exist", child);
      }
    }
    return true;
  }

  /**
   * Gets information of the given job id.
   *
   * @param jobId the id of the job
   * @param verbose whether the output should be verbose
   * @return null if the job id isn't know by the workflow tracker. WorkflowInfo otherwise
   */
  public WorkflowInfo getStatus(long jobId, boolean verbose) {
    WorkflowExecution workflowExecution = mWorkflows.get(jobId);

    if (workflowExecution == null) {
      return null;
    }

    ArrayList<Long> children = Lists.newArrayList(mChildren.get(jobId).iterator());
    Collections.sort(children);

    List<JobInfo> jobInfos = Lists.newArrayList();

    if (verbose) {
      for (long child : children) {
        try {
          jobInfos.add(mJobMaster.getStatus(child));
        } catch (JobDoesNotExistException e) {
          LOG.info(String.format("No job info on child job id %s. Skipping", child));
        }
      }
    }

    return new WorkflowInfo(jobId, workflowExecution.getName(),
        workflowExecution.getStatus(), workflowExecution.getLastUpdated(),
        workflowExecution.getErrorType(), workflowExecution.getErrorMessage(), jobInfos);
  }

  /**
   * @return all workflow infos
   */
  public Collection<WorkflowInfo> getAllInfo() {
    ArrayList<WorkflowInfo> res = Lists.newArrayList();

    for (Long workflowId : mWorkflows.keySet()) {
      res.add(getStatus(workflowId, false));
    }
    return res;
  }

  /**
   * @param name job name filter
   * @param statusList status list filter
   * @return job ids matching conditions
   */
  public Set<Long> findJobs(String name, List<Status> statusList) {
    Set<Long> jobs = new HashSet<>();
    for (Long workflowId : mWorkflows.entrySet().stream()
        .filter(x -> statusList.isEmpty() || statusList.contains(x.getValue().getStatus()))
        .map(Map.Entry::getKey).collect(Collectors.toList())) {
      if (name == null || name.isEmpty() || mWorkflows.get(workflowId).getName().equals(name)) {
        jobs.add(workflowId);
      }
    }
    return jobs;
  }

  /**
   * @return a list of all workflow ids
   */
  public Collection<Long> list() {
    return Collections.unmodifiableCollection(mWorkflows.keySet());
  }

  /**
   * Recursively cleanup the parent workflows given plans to be removed from the PlanTracker.
   * @param removedPlanIds the plan ids that are being cleaned up by PlanTracker
   */
  public synchronized void cleanup(Collection<Long> removedPlanIds) {
    for (long removedPlanId : removedPlanIds) {
      clean(removedPlanId);
    }
  }

  private synchronized void clean(long jobId) {
    mWorkflows.remove(jobId);
    mWaitingOn.remove(jobId);
    mChildren.remove(jobId);

    Long parentId = mParentWorkflow.remove(jobId);

    if (parentId == null) {
      return;
    }

    // TODO(bradley): share details of the child job to the parent workflow before deleting.
    ConcurrentHashSet<Long> siblings = mChildren.get(parentId);
    siblings.remove(jobId);

    if (siblings.isEmpty()) {
      WorkflowExecution parentExecution = mWorkflows.get(parentId);
      if (parentExecution != null && parentExecution.getStatus().isFinished()) {
        clean(parentId);
      }
    }
  }

  private synchronized void done(long jobId) {
    Long parentJobId = mParentWorkflow.get(jobId);

    if (parentJobId == null) {
      return;
    }

    ConcurrentHashSet<Long> siblings = mWaitingOn.get(parentJobId);

    siblings.remove(jobId);

    if (siblings.isEmpty()) {
      next(parentJobId);
    }
  }

  private synchronized void stop(long jobId, Status status, String errorType,
                                 String errorMessage) {
    Long parentJobId = mParentWorkflow.get(jobId);

    if (parentJobId == null) {
      return;
    }

    WorkflowExecution workflowExecution = mWorkflows.get(parentJobId);

    workflowExecution.stop(status, errorType, errorMessage);

    stop(parentJobId, status, errorType, errorMessage);
  }

  private synchronized void next(long jobId) {
    WorkflowExecution workflowExecution = mWorkflows.get(jobId);
    mChildren.putIfAbsent(jobId, new ConcurrentHashSet<>());

    Set<JobConfig> childJobConfigs = workflowExecution.next();

    if (childJobConfigs.isEmpty()) {
      done(jobId);
      return;
    }

    ConcurrentHashSet<Long> childJobIds = new ConcurrentHashSet<>();
    for (int i = 0; i < childJobConfigs.size(); i++) {
      childJobIds.add(mJobMaster.getNewJobId());
    }

    mWaitingOn.put(jobId, childJobIds);
    mChildren.get(jobId).addAll(childJobIds);

    for (Long childJobId : childJobIds) {
      mParentWorkflow.put(childJobId, jobId);
    }

    Iterator<Long> childJobIdsIter = childJobIds.iterator();
    Iterator<JobConfig> childJobConfigsIter = childJobConfigs.iterator();

    while (childJobIdsIter.hasNext() && childJobConfigsIter.hasNext()) {
      Long childJobId = childJobIdsIter.next();
      JobConfig childJobConfig = childJobConfigsIter.next();
      try {
        mJobMaster.run(childJobConfig, childJobId);
      } catch (JobDoesNotExistException | ResourceExhaustedException e) {
        LOG.warn(e.getMessage());
        final String errorType = ErrorUtils.getErrorType(e);
        workflowExecution.stop(Status.FAILED, errorType, e.getMessage());
        stop(jobId, Status.FAILED, errorType, e.getMessage());
      }
    }
  }

  /**
   * Updates internal state of the workflows based on the updated state of a plan.
   * @param planInfo info of the plan that had its status changed
   */
  public void onPlanStatusChange(PlanInfo planInfo) {
    Status status = planInfo.getStatus();
    switch (status) {
      case COMPLETED:
        done(planInfo.getId());
        break;
      case CANCELED:
      case FAILED:
        stop(planInfo.getId(), status, planInfo.getErrorType(), planInfo.getErrorMessage());
        break;
      default:
        break;
    }
  }
}
