package org.apache.helix.controller.rebalancer.util;

import org.apache.helix.HelixDataAccessor;
import org.apache.helix.HelixManager;
import org.apache.helix.PropertyKey;
import org.apache.helix.model.IdealState;
import org.apache.helix.task.TaskConstants;
import org.apache.helix.task.TaskUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

/**
 * Class for trigger rebalancing of a set of resource in a future time.
 */
public class RebalanceScheduler {
  private static final Logger LOG = LoggerFactory.getLogger(RebalanceScheduler.class);

  private class ScheduledTask {
    long _startTime;
    Future _future;

    public ScheduledTask(long _startTime, Future _future) {
      this._startTime = _startTime;
      this._future = _future;
    }

    public long getStartTime() {
      return _startTime;
    }

    public Future getFuture() {
      return _future;
    }
  }

  private final Map<String, ScheduledTask> _rebalanceTasks = new HashMap<String, ScheduledTask>();
  private final ScheduledExecutorService _rebalanceExecutor =
      Executors.newSingleThreadScheduledExecutor();

  /**
   * Add a future rebalance task for resource at given startTime
   *
   * @param resource
   * @param startTime time in milliseconds
   */
  public void scheduleRebalance(HelixManager manager, String resource, long startTime) {
    // Do nothing if there is already a timer set for the this workflow with the same start time.
    ScheduledTask existTask = _rebalanceTasks.get(resource);
    if (existTask != null && existTask.getStartTime() == startTime) {
      LOG.debug("Schedule timer for job: " + resource + " is up to date.");
      return;
    }

    long delay = startTime - System.currentTimeMillis();
    if (delay < 0) {
      LOG.debug(String.format("Delay time is %s, will not be scheduled", delay));
    }
    LOG.info(
        "Schedule rebalance for resource : " + resource + " at time: " + startTime + " delay: "
            + delay);

    // For workflow not yet scheduled, schedule them and record it
    RebalanceInvoker rebalanceInvoker = new RebalanceInvoker(manager, resource);
    ScheduledFuture future =
        _rebalanceExecutor.schedule(rebalanceInvoker, delay, TimeUnit.MILLISECONDS);
    ScheduledTask prevTask = _rebalanceTasks.put(resource, new ScheduledTask(startTime, future));
    if (prevTask != null && !prevTask.getFuture().isDone()) {
      if (!prevTask.getFuture().cancel(false)) {
        LOG.warn("Failed to cancel scheduled timer task for " + resource);
      }
      LOG.info("Remove previously scheduled timer task for " + resource);
    }
  }

  /**
   * Get the current schedule time for given resource.
   *
   * @param resource
   * @return existing schedule time or -1 if there is no scheduled task for this resource
   */
  public long getRebalanceTime(String resource) {
    ScheduledTask task = _rebalanceTasks.get(resource);
    if (task != null && !task.getFuture().isDone()) {
      return task.getStartTime();
    }
    return -1;
  }

  /**
   * Remove all existing future schedule tasks for the given resource
   *
   * @param resource
   */
  public long removeScheduledRebalance(String resource) {
    ScheduledTask existTask = _rebalanceTasks.remove(resource);
    if (existTask != null && !existTask.getFuture().isDone()) {
      if (!existTask.getFuture().cancel(true)) {
        LOG.warn("Failed to cancel scheduled timer task for " + resource);
      }
      LOG.info(
          "Remove scheduled rebalance task at time " + existTask.getStartTime() + " for resource: "
              + resource);

      return existTask.getStartTime();
    }

    return -1;
  }

  /**
   * The simplest possible runnable that will trigger a run of the controller pipeline
   */
  private class RebalanceInvoker implements Runnable {
    private final HelixManager _manager;
    private final String _resource;

    public RebalanceInvoker(HelixManager manager, String resource) {
      _manager = manager;
      _resource = resource;
    }

    @Override
    public void run() {
      invokeRebalance(_manager.getHelixDataAccessor(), _resource);
    }
  }

  /**
   * Trigger the controller to perform rebalance for a given resource.
   *
   * @param accessor Helix data accessor
   * @param resource the name of the resource changed to triggering the execution
   */
  public static void invokeRebalance(HelixDataAccessor accessor, String resource) {
    LOG.info("invoke rebalance for " + resource);
    PropertyKey key = accessor.keyBuilder().idealStates(resource);
    IdealState is = accessor.getProperty(key);
    if (is != null) {
      if (!accessor.updateProperty(key, is)) {
        LOG.warn("Failed to invoke rebalance on resource " + resource);
      }
    } else {
      LOG.warn("Can't find ideal state for " + resource);
    }
  }
}

@@ -211,90 +211,8 @@
 te;%0A
-import org.apache.helix.task.TaskConstants;%0Aimport org.apache.helix.task.TaskUtil;
 %0Aimp
@@ -1351,21 +1351,16 @@
 artTime%0A
-   *%0A
    * @pa
@@ -1780,32 +1780,18 @@
 or job: 
-%22 + resource + %22
+%7B%7D
  is up t
@@ -1798,16 +1798,26 @@
 o date.%22
+, resource
 );%0A     
@@ -2011,25 +2011,16 @@
 OG.info(
-%0A        
 %22Schedul
@@ -2050,68 +2050,63 @@
 e : 
-%22 + resource + %22 at time: %22 + startTime + %22 delay: %22
+%7B%7D at time: %7B%7D delay: %7B%7D%22, resource, startTime,
 %0A       
     
@@ -2101,22 +2101,16 @@
 %0A       
-     +
  delay);
@@ -2634,35 +2634,36 @@
  timer task for 
-%22 +
+%7B%7D%22,
  resource);%0A    
@@ -2717,35 +2717,36 @@
  timer task for 
-%22 +
+%7B%7D%22,
  resource);%0A    
@@ -2805,37 +2805,32 @@
 given resource.%0A
-   *%0A
    * @param reso
@@ -3226,21 +3226,16 @@
 esource%0A
-   *%0A
    * @pa
@@ -3586,27 +3586,16 @@
 OG.info(
-%0A          
 %22Remove 
@@ -3630,73 +3630,67 @@
 time
- %22 + existTask.getStartTime() + %22 for resource: %22%0A              +
+: %7B%7D for resource: %7B%7D%22,%0A          existTask.getStartTime(),
  res
@@ -4323,13 +4323,8 @@
 ce.%0A
-   *%0A
    *
@@ -4822,19 +4822,20 @@
 esource 
-%22 +
+%7B%7D%22,
  resourc
@@ -4898,27 +4898,28 @@
 l state for 
-%22 +
+%7B%7D%22,
  resource);%0A

Refactor log statements in RebalanceScheduler