package org.apache.helix.task;

/*
 * 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.
 */

import com.google.common.collect.Lists;
import org.I0Itec.zkclient.DataUpdater;
import org.apache.helix.*;
import org.apache.helix.controller.stages.ClusterDataCache;
import org.apache.helix.controller.stages.CurrentStateOutput;
import org.apache.helix.model.*;
import org.apache.helix.model.builder.CustomModeISBuilder;
import org.apache.helix.model.builder.IdealStateBuilder;
import org.apache.log4j.Logger;

import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TimeZone;

/**
 * Custom rebalancer implementation for the {@code Workflow} in task state model.
 */
public class WorkflowRebalancer extends TaskRebalancer {
  private static final Logger LOG = Logger.getLogger(WorkflowRebalancer.class);

  @Override
  public ResourceAssignment computeBestPossiblePartitionState(ClusterDataCache clusterData,
      IdealState taskIs, Resource resource, CurrentStateOutput currStateOutput) {
    final String workflow = resource.getResourceName();
    LOG.debug("Computer Best Partition for workflow: " + workflow);

    // Fetch workflow configuration and context
    WorkflowConfig workflowCfg = TaskUtil.getWorkflowCfg(_manager, workflow);
    if (workflowCfg == null) {
      LOG.warn("Workflow configuration is NULL for " + workflow);
      return buildEmptyAssignment(workflow, currStateOutput);
    }

    WorkflowContext workflowCtx = TaskUtil.getWorkflowContext(_manager, workflow);
    // Initialize workflow context if needed
    if (workflowCtx == null) {
      workflowCtx = new WorkflowContext(new ZNRecord(TaskUtil.WORKFLOW_CONTEXT_KW));
      workflowCtx.setStartTime(System.currentTimeMillis());
      LOG.debug("Workflow context is created for " + workflow);
    }

    // Clean up if workflow marked for deletion
    TargetState targetState = workflowCfg.getTargetState();
    if (targetState == TargetState.DELETE) {
      LOG.info("Workflow is marked as deleted " + workflow + " cleaning up the workflow context.");
      cleanupWorkflow(workflow, workflowCfg);
      return buildEmptyAssignment(workflow, currStateOutput);
    }

    if (targetState == TargetState.STOP) {
      LOG.info("Workflow " + workflow + "is marked as stopped.");
      return buildEmptyAssignment(workflow, currStateOutput);
    }

    long currentTime = System.currentTimeMillis();
    // Check if workflow has been finished and mark it if it is.
    if (workflowCtx.getFinishTime() == WorkflowContext.UNFINISHED
        && isWorkflowFinished(workflowCtx, workflowCfg)) {
      workflowCtx.setFinishTime(currentTime);
      TaskUtil.setWorkflowContext(_manager, workflow, workflowCtx);
    }

    if (workflowCtx.getFinishTime() != WorkflowContext.UNFINISHED) {
      LOG.info("Workflow " + workflow + " is finished.");
      long expiryTime = workflowCfg.getExpiry();
      // Check if this workflow has been finished past its expiry.
      if (workflowCtx.getFinishTime() + expiryTime <= currentTime) {
        LOG.info("Workflow " + workflow + " passed expiry time, cleaning up the workflow context.");
        cleanupWorkflow(workflow, workflowCfg);
      } else {
        // schedule future cleanup work
        long cleanupTime = workflowCtx.getFinishTime() + expiryTime;
        _scheduledRebalancer.scheduleRebalance(_manager, workflow, cleanupTime);
      }
      return buildEmptyAssignment(workflow, currStateOutput);
    }

    if (!isWorkflowReadyForSchedule(workflowCfg)) {
      LOG.info("Workflow " + workflow + " is not ready to schedule");
      // set the timer to trigger future schedule
      _scheduledRebalancer
          .scheduleRebalance(_manager, workflow, workflowCfg.getStartTime().getTime());
      return buildEmptyAssignment(workflow, currStateOutput);
    }

    // Check for readiness, and stop processing if it's not ready
    boolean isReady =
        scheduleWorkflowIfReady(workflow, workflowCfg, workflowCtx);
    if (isReady) {
      // Schedule jobs from this workflow.
      scheduleJobs(workflow, workflowCfg, workflowCtx);
    } else {
      LOG.debug("Workflow " + workflow + " is not ready to be scheduled.");
    }

    TaskUtil.setWorkflowContext(_manager, workflow, workflowCtx);
    return buildEmptyAssignment(workflow, currStateOutput);
  }

  /**
   * Figure out whether the jobs in the workflow should be run,
   * and if it's ready, then just schedule it
   */
  private void scheduleJobs(String workflow, WorkflowConfig workflowCfg, WorkflowContext workflowCtx) {
    ScheduleConfig scheduleConfig = workflowCfg.getScheduleConfig();
    if (scheduleConfig != null && scheduleConfig.isRecurring()) {
      LOG.debug("Jobs from recurring workflow are not schedule-able");
      return;
    }

    int scheduledJobs = 0;
    long timeToSchedule = Long.MAX_VALUE;
    for (String job : workflowCfg.getJobDag().getAllNodes()) {
      TaskState jobState = workflowCtx.getJobState(job);
      if (jobState != null && !jobState.equals(TaskState.NOT_STARTED)) {
        LOG.debug("Job " + job + " is already started or completed.");
        continue;
      }

      if (workflowCfg.isJobQueue() && scheduledJobs >= workflowCfg.getParallelJobs()) {
        LOG.debug(String.format("Workflow %s already have enough job in progress, "
                + "scheduledJobs(s)=%d, stop scheduling more jobs", workflow, scheduledJobs));
        break;
      }

      // check ancestor job status
      if (isJobReadyToSchedule(job, workflowCfg, workflowCtx)) {
        JobConfig jobConfig = TaskUtil.getJobCfg(_manager, job);
        // Since the start time is calculated base on the time of completion of parent jobs for this
        // job, the calculated start time should only be calculate once. Persist the calculated time
        // in WorkflowContext znode.
        Map<String, String> startTimeMap = workflowCtx.getRecord().getMapField(START_TIME_KEY);
        if (startTimeMap == null) {
          startTimeMap = new HashMap<String, String>();
          workflowCtx.getRecord().setMapField(START_TIME_KEY, startTimeMap);
        }

        long calculatedStartTime = System.currentTimeMillis();
        if (startTimeMap.containsKey(job)) {
          // Get the start time if it is already calculated
          calculatedStartTime = Long.parseLong(startTimeMap.get(job));
        } else {
          // If the start time is not calculated before, do the math.
          if (jobConfig.getExecutionDelay() >= 0) {
            calculatedStartTime += jobConfig.getExecutionDelay();
          }
          calculatedStartTime = Math.max(calculatedStartTime, jobConfig.getExecutionStart());
          startTimeMap.put(job, String.valueOf(calculatedStartTime));
          workflowCtx.getRecord().setMapField(START_TIME_KEY, startTimeMap);
          TaskUtil.setWorkflowContext(_manager, jobConfig.getWorkflow(), workflowCtx);
        }

        // Time is not ready. Set a trigger and update the start time.
        if (System.currentTimeMillis() < calculatedStartTime) {
          timeToSchedule = Math.min(timeToSchedule, calculatedStartTime);
        } else {
          scheduleSingleJob(job, jobConfig);
          scheduledJobs++;
        }
      }
    }
    if (timeToSchedule < Long.MAX_VALUE) {
      _scheduledRebalancer.scheduleRebalance(_manager, workflow, timeToSchedule);
    }
  }

  /**
   * Posts new job to cluster
   */
  private void scheduleSingleJob(String jobResource, JobConfig jobConfig) {
    HelixAdmin admin = _manager.getClusterManagmentTool();

    IdealState jobIS = admin.getResourceIdealState(_manager.getClusterName(), jobResource);
    if (jobIS != null) {
      LOG.info("Job " + jobResource + " idealstate already exists!");
      return;
    }

    // Set up job resource based on partitions from target resource
    TaskUtil.createUserContent(_manager.getHelixPropertyStore(), jobResource,
        new ZNRecord(TaskUtil.USER_CONTENT_NODE));
    int numIndependentTasks = jobConfig.getTaskConfigMap().size();

    int numPartitions = numIndependentTasks;
    if (numPartitions == 0) {
      IdealState targetIs =
          admin.getResourceIdealState(_manager.getClusterName(), jobConfig.getTargetResource());
      if (targetIs == null) {
        LOG.warn("Target resource does not exist for job " + jobResource);
        // do not need to fail here, the job will be marked as failure immediately when job starts running.
      } else {
        numPartitions = targetIs.getPartitionSet().size();
      }
    }

    admin.addResource(_manager.getClusterName(), jobResource, numPartitions,
        TaskConstants.STATE_MODEL_NAME);

    HelixDataAccessor accessor = _manager.getHelixDataAccessor();

    // Set the job configuration
    PropertyKey.Builder keyBuilder = accessor.keyBuilder();
    HelixProperty resourceConfig = new HelixProperty(jobResource);
    resourceConfig.getRecord().getSimpleFields().putAll(jobConfig.getResourceConfigMap());
    Map<String, TaskConfig> taskConfigMap = jobConfig.getTaskConfigMap();
    if (taskConfigMap != null) {
      for (TaskConfig taskConfig : taskConfigMap.values()) {
        resourceConfig.getRecord().setMapField(taskConfig.getId(), taskConfig.getConfigMap());
      }
    }
    accessor.setProperty(keyBuilder.resourceConfig(jobResource), resourceConfig);

    // Push out new ideal state based on number of target partitions
    IdealStateBuilder builder = new CustomModeISBuilder(jobResource);
    builder.setRebalancerMode(IdealState.RebalanceMode.TASK);
    builder.setNumReplica(1);
    builder.setNumPartitions(numPartitions);
    builder.setStateModel(TaskConstants.STATE_MODEL_NAME);

    if (jobConfig.getInstanceGroupTag() != null) {
      builder.setNodeGroup(jobConfig.getInstanceGroupTag());
    }

    if (jobConfig.isDisableExternalView()) {
      builder.setDisableExternalView(true);
    }

    jobIS = builder.build();
    for (int i = 0; i < numPartitions; i++) {
      jobIS.getRecord().setListField(jobResource + "_" + i, new ArrayList<String>());
      jobIS.getRecord().setMapField(jobResource + "_" + i, new HashMap<String, String>());
    }
    jobIS.setRebalancerClassName(JobRebalancer.class.getName());
    admin.setResourceIdealState(_manager.getClusterName(), jobResource, jobIS);
  }

  /**
   * Check if a workflow is ready to schedule, and schedule a rebalance if it is not
   *
   * @param workflow the Helix resource associated with the workflow
   * @param workflowCfg  the workflow to check
   * @param workflowCtx  the current workflow context
   * @return true if the workflow is ready for schedule, false if not ready
   */
  private boolean scheduleWorkflowIfReady(String workflow, WorkflowConfig workflowCfg,
      WorkflowContext workflowCtx) {
    // non-scheduled workflow is ready to run immediately.
    if (workflowCfg == null || workflowCfg.getScheduleConfig() == null) {
      return true;
    }

    // Figure out when this should be run, and if it's ready, then just run it
    ScheduleConfig scheduleConfig = workflowCfg.getScheduleConfig();
    Date startTime = scheduleConfig.getStartTime();
    long currentTime = new Date().getTime();
    long delayFromStart = startTime.getTime() - currentTime;

    if (delayFromStart <= 0) {
      // Recurring workflows are just templates that spawn new workflows
      if (scheduleConfig.isRecurring()) {
        // Skip scheduling this workflow if it's not in a start state
        if (!workflowCfg.getTargetState().equals(TargetState.START)) {
          LOG.debug("Skip scheduling since the workflow has not been started " + workflow);
          return false;
        }

        // Skip scheduling this workflow again if the previous run (if any) is still active
        String lastScheduled = workflowCtx.getLastScheduledSingleWorkflow();
        if (lastScheduled != null) {
          WorkflowContext lastWorkflowCtx = TaskUtil.getWorkflowContext(_manager, lastScheduled);
          if (lastWorkflowCtx != null
              && lastWorkflowCtx.getFinishTime() == WorkflowContext.UNFINISHED) {
            LOG.info("Skip scheduling since last schedule has not completed yet " + lastScheduled);
            return false;
          }
        }

        // Figure out how many jumps are needed, thus the time to schedule the next workflow
        // The negative of the delay is the amount of time past the start time
        long period =
            scheduleConfig.getRecurrenceUnit().toMillis(scheduleConfig.getRecurrenceInterval());
        long offsetMultiplier = (-delayFromStart) / period;
        long timeToSchedule = period * offsetMultiplier + startTime.getTime();

        // Now clone the workflow if this clone has not yet been created
        DateFormat df = new SimpleDateFormat("yyyyMMdd'T'HHmmss");
        df.setTimeZone(TimeZone.getTimeZone("UTC"));
        String newWorkflowName = workflow + "_" + df.format(new Date(timeToSchedule));
        LOG.debug("Ready to start workflow " + newWorkflowName);
        if (!newWorkflowName.equals(lastScheduled)) {
          Workflow clonedWf =
              cloneWorkflow(_manager, workflow, newWorkflowName, new Date(timeToSchedule));
          TaskDriver driver = new TaskDriver(_manager);
          try {
            // Start the cloned workflow
            driver.start(clonedWf);
          } catch (Exception e) {
            LOG.error("Failed to schedule cloned workflow " + newWorkflowName, e);
            _clusterStatusMonitor
                .updateWorkflowCounters(clonedWf.getWorkflowConfig(), TaskState.FAILED);
          }
          // Persist workflow start regardless of success to avoid retrying and failing
          workflowCtx.setLastScheduledSingleWorkflow(newWorkflowName);
          TaskUtil.setWorkflowContext(_manager, workflow, workflowCtx);
        }

        // Change the time to trigger the pipeline to that of the next run
        _scheduledRebalancer.scheduleRebalance(_manager, workflow, (timeToSchedule + period));
      } else {
        // one time workflow.
        // Remove any timers that are past-time for this workflowg
        long scheduledTime = _scheduledRebalancer.getRebalanceTime(workflow);
        if (scheduledTime > 0 && currentTime > scheduledTime) {
          _scheduledRebalancer.removeScheduledRebalance(workflow);
        }
        return true;
      }
    } else {
      // set the timer to trigger future schedule
      _scheduledRebalancer.scheduleRebalance(_manager, workflow, startTime.getTime());
    }

    return false;
  }

  /**
   * Create a new workflow based on an existing one
   *
   * @param manager          connection to Helix
   * @param origWorkflowName the name of the existing workflow
   * @param newWorkflowName  the name of the new workflow
   * @param newStartTime     a provided start time that deviates from the desired start time
   * @return the cloned workflow, or null if there was a problem cloning the existing one
   */
  public static Workflow cloneWorkflow(HelixManager manager, String origWorkflowName,
      String newWorkflowName, Date newStartTime) {
    // Read all resources, including the workflow and jobs of interest
    HelixDataAccessor accessor = manager.getHelixDataAccessor();
    PropertyKey.Builder keyBuilder = accessor.keyBuilder();
    Map<String, HelixProperty> resourceConfigMap =
        accessor.getChildValuesMap(keyBuilder.resourceConfigs());
    if (!resourceConfigMap.containsKey(origWorkflowName)) {
      LOG.error("No such workflow named " + origWorkflowName);
      return null;
    }
    if (resourceConfigMap.containsKey(newWorkflowName)) {
      LOG.error("Workflow with name " + newWorkflowName + " already exists!");
      return null;
    }

    // Create a new workflow with a new name
    Map<String, String> workflowConfigsMap =
        resourceConfigMap.get(origWorkflowName).getRecord().getSimpleFields();
    WorkflowConfig.Builder workflowConfigBlder = WorkflowConfig.Builder.fromMap(workflowConfigsMap);

    // Set the schedule, if applicable
    if (newStartTime != null) {
      ScheduleConfig scheduleConfig = ScheduleConfig.oneTimeDelayedStart(newStartTime);
      workflowConfigBlder.setScheduleConfig(scheduleConfig);
    }
    workflowConfigBlder.setTerminable(true);

    WorkflowConfig workflowConfig = workflowConfigBlder.build();

    JobDag jobDag = workflowConfig.getJobDag();
    Map<String, Set<String>> parentsToChildren = jobDag.getParentsToChildren();

    Workflow.Builder workflowBuilder = new Workflow.Builder(newWorkflowName);
    workflowBuilder.setWorkflowConfig(workflowConfig);

    // Add each job back as long as the original exists
    Set<String> namespacedJobs = jobDag.getAllNodes();
    for (String namespacedJob : namespacedJobs) {
      if (resourceConfigMap.containsKey(namespacedJob)) {
        // Copy over job-level and task-level configs
        String job = TaskUtil.getDenamespacedJobName(origWorkflowName, namespacedJob);
        HelixProperty jobConfig = resourceConfigMap.get(namespacedJob);
        Map<String, String> jobSimpleFields = jobConfig.getRecord().getSimpleFields();

        JobConfig.Builder jobCfgBuilder = JobConfig.Builder.fromMap(jobSimpleFields);

        jobCfgBuilder.setWorkflow(newWorkflowName); // overwrite workflow name
        Map<String, Map<String, String>> rawTaskConfigMap = jobConfig.getRecord().getMapFields();
        List<TaskConfig> taskConfigs = Lists.newLinkedList();
        for (Map<String, String> rawTaskConfig : rawTaskConfigMap.values()) {
          TaskConfig taskConfig = TaskConfig.Builder.from(rawTaskConfig);
          taskConfigs.add(taskConfig);
        }
        jobCfgBuilder.addTaskConfigs(taskConfigs);
        workflowBuilder.addJob(job, jobCfgBuilder);

        // Add dag dependencies
        Set<String> children = parentsToChildren.get(namespacedJob);
        if (children != null) {
          for (String namespacedChild : children) {
            String child = TaskUtil.getDenamespacedJobName(origWorkflowName, namespacedChild);
            workflowBuilder.addParentChildDependency(job, child);
          }
        }
      }
    }
    return workflowBuilder.build();
  }

  /**
   * Cleans up workflow configs and workflow contexts associated with this workflow,
   * including all job-level configs and context, plus workflow-level information.
   */
  private void cleanupWorkflow(String workflow, WorkflowConfig workflowcfg) {
    LOG.info("Cleaning up workflow: " + workflow);
    HelixDataAccessor accessor = _manager.getHelixDataAccessor();

    /*
    if (workflowCtx != null && workflowCtx.getFinishTime() == WorkflowContext.UNFINISHED) {
      LOG.error("Workflow " + workflow + " has not completed, abort the clean up task.");
      return;
    }*/

    for (String job : workflowcfg.getJobDag().getAllNodes()) {
      cleanupJob(job, workflow);
    }

    // clean up workflow-level info if this was the last in workflow
    if (workflowcfg.isTerminable() || workflowcfg.getTargetState() == TargetState.DELETE) {
      // clean up IS & EV
      cleanupIdealStateExtView(_manager.getHelixDataAccessor(), workflow);

      // delete workflow config
      PropertyKey workflowCfgKey = TaskUtil.getWorkflowConfigKey(accessor, workflow);
      if (accessor.getProperty(workflowCfgKey) != null) {
        if (!accessor.removeProperty(workflowCfgKey)) {
          LOG.error(String.format(
              "Error occurred while trying to clean up workflow %s. Failed to remove node %s from Helix.",
              workflow, workflowCfgKey));
        }
      }
      // Delete workflow context
      LOG.info("Removing workflow context: " + workflow);
      if (!TaskUtil.removeWorkflowContext(_manager, workflow)) {
        LOG.error(String.format(
            "Error occurred while trying to clean up workflow %s. Aborting further clean up steps.",
            workflow));
      }

      // Remove pending timer task for this workflow if exists
      _scheduledRebalancer.removeScheduledRebalance(workflow);
    }
  }


  /**
   * Cleans up workflow configs and workflow contexts associated with this workflow,
   * including all job-level configs and context, plus workflow-level information.
   */
  private void cleanupJob(final String job, String workflow) {
    LOG.info("Cleaning up job: " + job + " in workflow: " + workflow);
    HelixDataAccessor accessor = _manager.getHelixDataAccessor();

    // Remove any idealstate and externalView.
    cleanupIdealStateExtView(accessor, job);

    // Remove DAG references in workflow
    PropertyKey workflowKey = TaskUtil.getWorkflowConfigKey(accessor, workflow);
    DataUpdater<ZNRecord> dagRemover = new DataUpdater<ZNRecord>() {
      @Override public ZNRecord update(ZNRecord currentData) {
        if (currentData != null) {
          JobDag jobDag = JobDag.fromJson(
              currentData.getSimpleField(WorkflowConfig.WorkflowConfigProperty.Dag.name()));
          for (String child : jobDag.getDirectChildren(job)) {
            jobDag.getChildrenToParents().get(child).remove(job);
          }
          for (String parent : jobDag.getDirectParents(job)) {
            jobDag.getParentsToChildren().get(parent).remove(job);
          }
          jobDag.getChildrenToParents().remove(job);
          jobDag.getParentsToChildren().remove(job);
          jobDag.getAllNodes().remove(job);
          try {
            currentData
                .setSimpleField(WorkflowConfig.WorkflowConfigProperty.Dag.name(), jobDag.toJson());
          } catch (Exception e) {
            LOG.error("Could not update DAG for job: " + job, e);
          }
        } else {
          LOG.error("Could not update DAG for job: " + job + " ZNRecord is null.");
        }
        return currentData;
      }
    };
    accessor.getBaseDataAccessor().update(workflowKey.getPath(), dagRemover,
        AccessOption.PERSISTENT);

    // Delete job configs.
    PropertyKey cfgKey = TaskUtil.getWorkflowConfigKey(accessor, job);
    if (accessor.getProperty(cfgKey) != null) {
      if (!accessor.removeProperty(cfgKey)) {
        LOG.error(String.format(
            "Error occurred while trying to clean up job %s. Failed to remove node %s from Helix.",
            job, cfgKey));
      }
    }

    // Delete job context
    // For recurring workflow, it's OK if the node doesn't exist.
    if (!TaskUtil.removeJobContext(_manager, job)) {
      LOG.warn(String.format("Error occurred while trying to clean up job %s.", job));
    }

    LOG.info(String.format("Successfully cleaned up job context %s.", job));

    _scheduledRebalancer.removeScheduledRebalance(job);
  }

  @Override
  public IdealState computeNewIdealState(String resourceName, IdealState currentIdealState,
      CurrentStateOutput currentStateOutput, ClusterDataCache clusterData) {
    // Nothing to do here with workflow resource.
    return currentIdealState;
  }
}
