/**
 * 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 org.apache.aurora.scheduler.thrift;

import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;

import javax.annotation.Nullable;
import javax.inject.Inject;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ContiguousSet;
import com.google.common.collect.DiscreteDomain;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Range;
import com.google.common.collect.Sets;

import org.apache.aurora.common.stats.StatsProvider;
import org.apache.aurora.gen.DrainHostsResult;
import org.apache.aurora.gen.EndMaintenanceResult;
import org.apache.aurora.gen.ExplicitReconciliationSettings;
import org.apache.aurora.gen.Hosts;
import org.apache.aurora.gen.InstanceKey;
import org.apache.aurora.gen.InstanceTaskConfig;
import org.apache.aurora.gen.JobConfiguration;
import org.apache.aurora.gen.JobKey;
import org.apache.aurora.gen.JobUpdate;
import org.apache.aurora.gen.JobUpdateInstructions;
import org.apache.aurora.gen.JobUpdateKey;
import org.apache.aurora.gen.JobUpdatePulseStatus;
import org.apache.aurora.gen.JobUpdateQuery;
import org.apache.aurora.gen.JobUpdateRequest;
import org.apache.aurora.gen.JobUpdateSettings;
import org.apache.aurora.gen.JobUpdateSummary;
import org.apache.aurora.gen.ListBackupsResult;
import org.apache.aurora.gen.MaintenanceStatusResult;
import org.apache.aurora.gen.PulseJobUpdateResult;
import org.apache.aurora.gen.QueryRecoveryResult;
import org.apache.aurora.gen.ReadOnlyScheduler;
import org.apache.aurora.gen.ResourceAggregate;
import org.apache.aurora.gen.Response;
import org.apache.aurora.gen.Result;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.gen.SlaPolicy;
import org.apache.aurora.gen.StartJobUpdateResult;
import org.apache.aurora.gen.StartMaintenanceResult;
import org.apache.aurora.gen.TaskQuery;
import org.apache.aurora.gen.VariableBatchJobUpdateStrategy;
import org.apache.aurora.scheduler.base.JobKeys;
import org.apache.aurora.scheduler.base.Numbers;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.configuration.ConfigurationManager;
import org.apache.aurora.scheduler.configuration.ConfigurationManager.TaskDescriptionException;
import org.apache.aurora.scheduler.configuration.SanitizedConfiguration;
import org.apache.aurora.scheduler.cron.CronException;
import org.apache.aurora.scheduler.cron.CronJobManager;
import org.apache.aurora.scheduler.cron.SanitizedCronJob;
import org.apache.aurora.scheduler.maintenance.MaintenanceController;
import org.apache.aurora.scheduler.quota.QuotaCheckResult;
import org.apache.aurora.scheduler.quota.QuotaManager;
import org.apache.aurora.scheduler.quota.QuotaManager.QuotaException;
import org.apache.aurora.scheduler.reconciliation.TaskReconciler;
import org.apache.aurora.scheduler.state.StateChangeResult;
import org.apache.aurora.scheduler.state.StateManager;
import org.apache.aurora.scheduler.state.UUIDGenerator;
import org.apache.aurora.scheduler.storage.SnapshotStore;
import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
import org.apache.aurora.scheduler.storage.Storage.NonVolatileStorage;
import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
import org.apache.aurora.scheduler.storage.backup.Recovery;
import org.apache.aurora.scheduler.storage.backup.StorageBackup;
import org.apache.aurora.scheduler.storage.durability.ThriftBackfill;
import org.apache.aurora.scheduler.storage.entities.IHostStatus;
import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IJobUpdate;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateKey;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateRequest;
import org.apache.aurora.scheduler.storage.entities.IJobUpdateSettings;
import org.apache.aurora.scheduler.storage.entities.IMetadata;
import org.apache.aurora.scheduler.storage.entities.IRange;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.apache.aurora.scheduler.thrift.aop.AnnotatedAuroraAdmin;
import org.apache.aurora.scheduler.thrift.aop.ThriftWorkload;
import org.apache.aurora.scheduler.thrift.auth.DecoratedThrift;
import org.apache.aurora.scheduler.updater.JobDiff;
import org.apache.aurora.scheduler.updater.JobUpdateController;
import org.apache.aurora.scheduler.updater.JobUpdateController.AuditData;
import org.apache.aurora.scheduler.updater.JobUpdateController.JobUpdatingException;
import org.apache.aurora.scheduler.updater.UpdateInProgressException;
import org.apache.aurora.scheduler.updater.UpdateStateException;
import org.apache.thrift.TException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

import static org.apache.aurora.common.base.MorePreconditions.checkNotBlank;
import static org.apache.aurora.gen.ResponseCode.INVALID_REQUEST;
import static org.apache.aurora.gen.ResponseCode.JOB_UPDATING_ERROR;
import static org.apache.aurora.gen.ResponseCode.OK;
import static org.apache.aurora.scheduler.base.Numbers.convertRanges;
import static org.apache.aurora.scheduler.base.Numbers.toRanges;
import static org.apache.aurora.scheduler.base.Tasks.ACTIVE_STATES;
import static org.apache.aurora.scheduler.base.Tasks.TERMINAL_STATES;
import static org.apache.aurora.scheduler.quota.QuotaCheckResult.Result.INSUFFICIENT_QUOTA;
import static org.apache.aurora.scheduler.thrift.Responses.addMessage;
import static org.apache.aurora.scheduler.thrift.Responses.empty;
import static org.apache.aurora.scheduler.thrift.Responses.error;
import static org.apache.aurora.scheduler.thrift.Responses.invalidRequest;
import static org.apache.aurora.scheduler.thrift.Responses.ok;

/**
 * Aurora scheduler thrift server implementation.
 * <p/>
 * Interfaces between users and the scheduler to access/modify jobs and perform cluster
 * administration tasks.
 */
@DecoratedThrift
class SchedulerThriftInterface implements AnnotatedAuroraAdmin {

  @VisibleForTesting
  static final String STAT_PREFIX = "thrift_workload_";
  @VisibleForTesting
  static final String CREATE_JOB = STAT_PREFIX + "createJob";
  @VisibleForTesting
  static final String CREATE_OR_UPDATE_CRON = STAT_PREFIX + "createOrUpdateCronTemplate";
  @VisibleForTesting
  static final String KILL_TASKS = STAT_PREFIX + "killTasks";
  @VisibleForTesting
  static final String RESTART_SHARDS = STAT_PREFIX + "restartShards";
  @VisibleForTesting
  static final String START_MAINTENANCE = STAT_PREFIX + "startMaintenance";
  @VisibleForTesting
  static final String DRAIN_HOSTS = STAT_PREFIX + "drainHosts";
  @VisibleForTesting
  static final String SLA_DRAIN_HOSTS = STAT_PREFIX + "slaDrainHosts";
  @VisibleForTesting
  static final String MAINTENANCE_STATUS = STAT_PREFIX + "maintenanceStatus";
  @VisibleForTesting
  static final String END_MAINTENANCE = STAT_PREFIX + "endMaintenance";
  @VisibleForTesting
  static final String ADD_INSTANCES = STAT_PREFIX + "addInstances";
  @VisibleForTesting
  static final String START_JOB_UPDATE = STAT_PREFIX + "startJobUpdate";
  @VisibleForTesting
  static final String PRUNE_TASKS = STAT_PREFIX + "pruneTasks";

  private static final Logger LOG = LoggerFactory.getLogger(SchedulerThriftInterface.class);

  private final ConfigurationManager configurationManager;
  private final Thresholds thresholds;
  private final NonVolatileStorage storage;
  private final SnapshotStore snapshotStore;
  private final StorageBackup backup;
  private final Recovery recovery;
  private final MaintenanceController maintenance;
  private final CronJobManager cronJobManager;
  private final QuotaManager quotaManager;
  private final StateManager stateManager;
  private final UUIDGenerator uuidGenerator;
  private final JobUpdateController jobUpdateController;
  private final ReadOnlyScheduler.Iface readOnlyScheduler;
  private final AuditMessages auditMessages;
  private final TaskReconciler taskReconciler;

  private final AtomicLong createJobCounter;
  private final AtomicLong createOrUpdateCronCounter;
  private final AtomicLong killTasksCounter;
  private final AtomicLong restartShardsCounter;
  private final AtomicLong startMaintenanceCounter;
  private final AtomicLong drainHostsCounter;
  private final AtomicLong slaDrainHostsCounter;
  private final AtomicLong maintenanceStatusCounter;
  private final AtomicLong endMaintenanceCounter;
  private final AtomicLong addInstancesCounter;
  private final AtomicLong startJobUpdateCounter;
  private final AtomicLong pruneTasksCounter;

  @Inject
  SchedulerThriftInterface(
      ConfigurationManager configurationManager,
      Thresholds thresholds,
      NonVolatileStorage storage,
      SnapshotStore snapshotStore,
      StorageBackup backup,
      Recovery recovery,
      CronJobManager cronJobManager,
      MaintenanceController maintenance,
      QuotaManager quotaManager,
      StateManager stateManager,
      UUIDGenerator uuidGenerator,
      JobUpdateController jobUpdateController,
      ReadOnlyScheduler.Iface readOnlyScheduler,
      AuditMessages auditMessages,
      TaskReconciler taskReconciler,
      StatsProvider statsProvider) {

    this.configurationManager = requireNonNull(configurationManager);
    this.thresholds = requireNonNull(thresholds);
    this.storage = requireNonNull(storage);
    this.snapshotStore = requireNonNull(snapshotStore);
    this.backup = requireNonNull(backup);
    this.recovery = requireNonNull(recovery);
    this.maintenance = requireNonNull(maintenance);
    this.cronJobManager = requireNonNull(cronJobManager);
    this.quotaManager = requireNonNull(quotaManager);
    this.stateManager = requireNonNull(stateManager);
    this.uuidGenerator = requireNonNull(uuidGenerator);
    this.jobUpdateController = requireNonNull(jobUpdateController);
    this.readOnlyScheduler = requireNonNull(readOnlyScheduler);
    this.auditMessages = requireNonNull(auditMessages);
    this.taskReconciler = requireNonNull(taskReconciler);

    this.createJobCounter = statsProvider.makeCounter(CREATE_JOB);
    this.createOrUpdateCronCounter = statsProvider.makeCounter(CREATE_OR_UPDATE_CRON);
    this.killTasksCounter = statsProvider.makeCounter(KILL_TASKS);
    this.restartShardsCounter = statsProvider.makeCounter(RESTART_SHARDS);
    this.startMaintenanceCounter = statsProvider.makeCounter(START_MAINTENANCE);
    this.drainHostsCounter = statsProvider.makeCounter(DRAIN_HOSTS);
    this.slaDrainHostsCounter = statsProvider.makeCounter(SLA_DRAIN_HOSTS);
    this.maintenanceStatusCounter = statsProvider.makeCounter(MAINTENANCE_STATUS);
    this.endMaintenanceCounter = statsProvider.makeCounter(END_MAINTENANCE);
    this.addInstancesCounter = statsProvider.makeCounter(ADD_INSTANCES);
    this.startJobUpdateCounter = statsProvider.makeCounter(START_JOB_UPDATE);
    this.pruneTasksCounter = statsProvider.makeCounter(PRUNE_TASKS);
  }

  @Override
  public Response createJob(JobConfiguration mutableJob) {
    SanitizedConfiguration sanitized;
    try {
      sanitized = SanitizedConfiguration.fromUnsanitized(
          configurationManager,
          IJobConfiguration.build(mutableJob));
    } catch (TaskDescriptionException e) {
      return error(INVALID_REQUEST, e);
    }

    if (sanitized.isCron()) {
      return invalidRequest(NO_CRON);
    }

    return storage.write(storeProvider -> {
      IJobConfiguration job = sanitized.getJobConfig();

      try {
        jobUpdateController.assertNotUpdating(job.getKey());

        checkJobExists(storeProvider, job.getKey());

        ITaskConfig template = sanitized.getJobConfig().getTaskConfig();
        int count = sanitized.getJobConfig().getInstanceCount();

        validateTaskLimits(
            count,
            quotaManager.checkInstanceAddition(template, count, storeProvider));

        LOG.info("Launching " + count + " tasks.");
        stateManager.insertPendingTasks(
            storeProvider,
            template,
            sanitized.getInstanceIds());
        createJobCounter.addAndGet(sanitized.getInstanceIds().size());

        return ok();
      } catch (JobUpdatingException e) {
        return error(JOB_UPDATING_ERROR, e);
      } catch (JobExistsException | TaskValidationException e) {
        return error(INVALID_REQUEST, e);
      }
    });
  }

  private static class JobExistsException extends Exception {
    JobExistsException(String message) {
      super(message);
    }
  }

  private void checkJobExists(StoreProvider store, IJobKey jobKey) throws JobExistsException {
    if (!Iterables.isEmpty(store.getTaskStore().fetchTasks(Query.jobScoped(jobKey).active()))
        || getCronJob(store, jobKey).isPresent()) {

      throw new JobExistsException(jobAlreadyExistsMessage(jobKey));
    }
  }

  private Response createOrUpdateCronTemplate(
      JobConfiguration mutableJob,
      boolean updateOnly) {

    IJobConfiguration job = IJobConfiguration.build(mutableJob);
    IJobKey jobKey = JobKeys.assertValid(job.getKey());

    SanitizedConfiguration sanitized;
    try {
      sanitized = SanitizedConfiguration.fromUnsanitized(configurationManager, job);
    } catch (TaskDescriptionException e) {
      return error(INVALID_REQUEST, e);
    }

    if (!sanitized.isCron()) {
      return invalidRequest(noCronScheduleMessage(jobKey));
    }

    return storage.write(storeProvider -> {
      try {
        jobUpdateController.assertNotUpdating(jobKey);

        int count = sanitized.getJobConfig().getInstanceCount();

        validateTaskLimits(
            count,
            quotaManager.checkCronUpdate(sanitized.getJobConfig(), storeProvider));

        // TODO(mchucarroll): Merge CronJobManager.createJob/updateJob
        if (updateOnly || getCronJob(storeProvider, jobKey).isPresent()) {
          // The job already has a schedule: so update it.
          cronJobManager.updateJob(SanitizedCronJob.from(sanitized));
        } else {
          checkJobExists(storeProvider, jobKey);
          cronJobManager.createJob(SanitizedCronJob.from(sanitized));
        }
        createOrUpdateCronCounter.addAndGet(count);

        return ok();
      } catch (JobUpdatingException e) {
        return error(JOB_UPDATING_ERROR, e);
      } catch (JobExistsException | TaskValidationException | CronException e) {
        return error(INVALID_REQUEST, e);
      }
    });
  }

  @Override
  public Response scheduleCronJob(JobConfiguration mutableJob) {
    return createOrUpdateCronTemplate(mutableJob, false);
  }

  @Override
  public Response replaceCronTemplate(JobConfiguration mutableJob) {
    return createOrUpdateCronTemplate(mutableJob, true);
  }

  @Override
  public Response descheduleCronJob(JobKey mutableJobKey) {
    try {
      IJobKey jobKey = JobKeys.assertValid(IJobKey.build(mutableJobKey));
      jobUpdateController.assertNotUpdating(jobKey);

      if (cronJobManager.deleteJob(jobKey)) {
        return ok();
      } else {
        return addMessage(empty(), OK, notScheduledCronMessage(jobKey));
      }
    } catch (JobUpdatingException e) {
      return error(JOB_UPDATING_ERROR, e);
    }
  }

  @Override
  public Response populateJobConfig(JobConfiguration description) throws TException {
    return readOnlyScheduler.populateJobConfig(description);
  }

  @Override
  public Response startCronJob(JobKey mutableJobKey) {
    IJobKey jobKey = JobKeys.assertValid(IJobKey.build(mutableJobKey));

    try {
      cronJobManager.startJobNow(jobKey);
      return ok();
    } catch (CronException e) {
      return invalidRequest("Failed to start cron job - " + e.getMessage());
    }
  }

  // TODO(William Farner): Provide status information about cron jobs here.
  @ThriftWorkload
  @Override
  public Response getTasksStatus(TaskQuery query) throws TException {
    return readOnlyScheduler.getTasksStatus(query);
  }

  @ThriftWorkload
  @Override
  public Response getTasksWithoutConfigs(TaskQuery query) throws TException {
    return readOnlyScheduler.getTasksWithoutConfigs(query);
  }

  @ThriftWorkload
  @Override
  public Response getPendingReason(TaskQuery query) throws TException {
    return readOnlyScheduler.getPendingReason(query);
  }

  @ThriftWorkload
  @Override
  public Response getConfigSummary(JobKey job) throws TException {
    return readOnlyScheduler.getConfigSummary(job);
  }

  @ThriftWorkload
  @Override
  public Response getRoleSummary() throws TException {
    return readOnlyScheduler.getRoleSummary();
  }

  @ThriftWorkload
  @Override
  public Response getJobSummary(@Nullable String maybeNullRole) throws TException {
    return readOnlyScheduler.getJobSummary(maybeNullRole);
  }

  @ThriftWorkload
  @Override
  public Response getJobs(@Nullable String maybeNullRole) throws TException {
    return readOnlyScheduler.getJobs(maybeNullRole);
  }

  @Override
  public Response getJobUpdateDiff(JobUpdateRequest request) throws TException {
    return readOnlyScheduler.getJobUpdateDiff(request);
  }

  @Override
  public Response getTierConfigs() throws TException {
    return readOnlyScheduler.getTierConfigs();
  }

  private static Query.Builder implicitKillQuery(Query.Builder query) {
    // Unless statuses were specifically supplied, only attempt to kill active tasks.
    return query.get().getStatuses().isEmpty() ? query.byStatus(ACTIVE_STATES) : query;
  }

  @Override
  public Response killTasks(
      @Nullable JobKey mutableJob,
      @Nullable Set<Integer> instances,
      @Nullable String message) {

    Response response = empty();
    IJobKey jobKey = JobKeys.assertValid(IJobKey.build(mutableJob));
    Query.Builder query;
    if (instances == null || Iterables.isEmpty(instances)) {
      query = implicitKillQuery(Query.jobScoped(jobKey));
    } else {
      query = implicitKillQuery(Query.instanceScoped(jobKey, instances));
    }

    return storage.write(storeProvider -> {
      try {
        jobUpdateController.assertNotUpdating(jobKey);
      } catch (JobUpdatingException e) {
        return error(JOB_UPDATING_ERROR, e);
      }

      Iterable<IScheduledTask> tasks = storeProvider.getTaskStore().fetchTasks(query);

      LOG.info("Killing tasks matching " + query);

      int tasksKilled = 0;
      for (String taskId : Tasks.ids(tasks)) {
        if (StateChangeResult.SUCCESS == stateManager.changeState(
            storeProvider,
            taskId,
            Optional.empty(),
            ScheduleStatus.KILLING,
            auditMessages.killedByRemoteUser(Optional.ofNullable(message)))) {
          ++tasksKilled;
        }
      }
      killTasksCounter.addAndGet(tasksKilled);

      return tasksKilled > 0
          ? response.setResponseCode(OK)
          : addMessage(response, OK, NO_TASKS_TO_KILL_MESSAGE);
    });
  }

  @Override
  public Response restartShards(JobKey mutableJobKey, Set<Integer> shardIds) {
    IJobKey jobKey = JobKeys.assertValid(IJobKey.build(mutableJobKey));
    checkNotBlank(shardIds);

    return storage.write(storeProvider -> {
      try {
        jobUpdateController.assertNotUpdating(jobKey);
      } catch (JobUpdatingException e) {
        return error(JOB_UPDATING_ERROR, e);
      }

      Query.Builder query = Query.instanceScoped(jobKey, shardIds).active();
      Iterable<IScheduledTask> matchingTasks = storeProvider.getTaskStore().fetchTasks(query);
      if (Iterables.size(matchingTasks) != shardIds.size()) {
        return invalidRequest("Not all requested shards are active.");
      }

      LOG.info("Restarting shards matching " + query);
      for (String taskId : Tasks.ids(matchingTasks)) {
        stateManager.changeState(
            storeProvider,
            taskId,
            Optional.empty(),
            ScheduleStatus.RESTARTING,
            auditMessages.restartedByRemoteUser());
      }
      restartShardsCounter.addAndGet(shardIds.size());

      return ok();
    });
  }

  @Override
  public Response getQuota(String ownerRole) throws TException {
    return readOnlyScheduler.getQuota(ownerRole);
  }

  @Override
  public Response setQuota(String ownerRole, ResourceAggregate resourceAggregate) {
    checkNotBlank(ownerRole);
    requireNonNull(resourceAggregate);

    try {
      storage.write((NoResult<QuotaException>) store -> quotaManager.saveQuota(
          ownerRole,
          ThriftBackfill.backfillResourceAggregate(resourceAggregate),
          store));
      return ok();
    } catch (QuotaException e) {
      return error(INVALID_REQUEST, e);
    }
  }

  @Override
  public Response startMaintenance(Hosts hosts) {
    startMaintenanceCounter.addAndGet(hosts.getHostNamesSize());
    return ok(Result.startMaintenanceResult(
        new StartMaintenanceResult()
            .setStatuses(IHostStatus.toBuildersSet(
                maintenance.startMaintenance(hosts.getHostNames())))));
  }

  @Override
  public Response drainHosts(Hosts hosts) {
    drainHostsCounter.addAndGet(hosts.getHostNamesSize());
    return ok(Result.drainHostsResult(
        new DrainHostsResult().setStatuses(IHostStatus.toBuildersSet(
            maintenance.drain(hosts.getHostNames())))));
  }

  @Override
  public Response maintenanceStatus(Hosts hosts) {
    maintenanceStatusCounter.addAndGet(hosts.getHostNamesSize());
    return ok(Result.maintenanceStatusResult(
        new MaintenanceStatusResult().setStatuses(IHostStatus.toBuildersSet(
            maintenance.getStatus(hosts.getHostNames())))));
  }

  @Override
  public Response endMaintenance(Hosts hosts) {
    endMaintenanceCounter.addAndGet(hosts.getHostNamesSize());
    return ok(Result.endMaintenanceResult(
        new EndMaintenanceResult()
            .setStatuses(IHostStatus.toBuildersSet(
                maintenance.endMaintenance(hosts.getHostNames())))));
  }

  @Override
  public Response slaDrainHosts(Hosts hosts, SlaPolicy defaultSlaPolicy, long timeoutSecs) {
    slaDrainHostsCounter.addAndGet(hosts.getHostNamesSize());
    return ok(Result.drainHostsResult(
        new DrainHostsResult().setStatuses(IHostStatus.toBuildersSet(
            maintenance.slaDrain(hosts.getHostNames(), defaultSlaPolicy, timeoutSecs)))));
  }

  @Override
  public Response forceTaskState(String taskId, ScheduleStatus status) {
    checkNotBlank(taskId);
    requireNonNull(status);

    storage.write(storeProvider -> stateManager.changeState(
        storeProvider,
        taskId,
        Optional.empty(),
        status,
        auditMessages.transitionedBy()));

    return ok();
  }

  @Override
  public Response performBackup() {
    backup.backupNow();
    return ok();
  }

  @Override
  public Response listBackups() {
    return ok(Result.listBackupsResult(new ListBackupsResult()
        .setBackups(recovery.listBackups())));
  }

  @Override
  public Response stageRecovery(String backupId) {
    recovery.stage(backupId);
    return ok();
  }

  @Override
  public Response queryRecovery(TaskQuery query) {
    return ok(Result.queryRecoveryResult(new QueryRecoveryResult()
        .setTasks(IScheduledTask.toBuildersSet(recovery.query(Query.arbitrary(query))))));
  }

  @Override
  public Response deleteRecoveryTasks(TaskQuery query) {
    recovery.deleteTasks(Query.arbitrary(query));
    return ok();
  }

  @Override
  public Response commitRecovery() {
    recovery.commit();
    return ok();
  }

  @Override
  public Response unloadRecovery() {
    recovery.unload();
    return ok();
  }

  @Override
  public Response snapshot() {
    snapshotStore.snapshot();
    return ok();
  }

  @Override
  public Response triggerExplicitTaskReconciliation(ExplicitReconciliationSettings settings)
      throws TException {
    try {
      requireNonNull(settings);
      Preconditions.checkArgument(!settings.isSetBatchSize() || settings.getBatchSize() > 0,
          "Batch size must be greater than zero.");

      Optional<Integer> batchSize = settings.isSetBatchSize()
          ? Optional.of(settings.getBatchSize())
          : Optional.empty();

      taskReconciler.triggerExplicitReconciliation(batchSize);
      return ok();
    } catch (IllegalArgumentException e) {
      return error(INVALID_REQUEST, e);
    }
  }

  @Override
  public Response triggerImplicitTaskReconciliation() throws TException {
    taskReconciler.triggerImplicitReconciliation();
    return ok();
  }

  @Override
  public Response addInstances(InstanceKey key, int count) {
    IJobKey jobKey = JobKeys.assertValid(IJobKey.build(key.getJobKey()));

    if (count <= 0) {
      return invalidRequest(INVALID_INSTANCE_COUNT);
    }

    Response response = empty();
    return storage.write(storeProvider -> {
      try {
        if (getCronJob(storeProvider, jobKey).isPresent()) {
          return invalidRequest("Instances may not be added to cron jobs.");
        }

        jobUpdateController.assertNotUpdating(jobKey);

        List<IScheduledTask> currentTasks = ImmutableList.copyOf(
            storeProvider.getTaskStore().fetchTasks(Query.jobScoped(jobKey).active()));

        Optional<IScheduledTask> templateTask = currentTasks.stream()
            .filter(e -> e.getAssignedTask().getInstanceId() == key.getInstanceId())
            .findFirst();
        if (!templateTask.isPresent()) {
          return invalidRequest(INVALID_INSTANCE_ID);
        }

        int lastId = currentTasks.stream()
            .map(e -> e.getAssignedTask().getInstanceId())
            .max(Comparator.naturalOrder())
            .get();

        Set<Integer> instanceIds = ContiguousSet.create(
            Range.openClosed(lastId, lastId + count),
            DiscreteDomain.integers());

        ITaskConfig task = templateTask.get().getAssignedTask().getTask();
        validateTaskLimits(
            Iterables.size(currentTasks) + instanceIds.size(),
            quotaManager.checkInstanceAddition(task, instanceIds.size(), storeProvider));

        stateManager.insertPendingTasks(storeProvider, task, instanceIds);
        addInstancesCounter.addAndGet(instanceIds.size());

        return response.setResponseCode(OK);
      } catch (JobUpdatingException e) {
        return error(JOB_UPDATING_ERROR, e);
      } catch (TaskValidationException | IllegalArgumentException e) {
        return error(INVALID_REQUEST, e);
      }
    });
  }

  private Optional<IJobConfiguration> getCronJob(StoreProvider storeProvider, IJobKey jobKey) {
    requireNonNull(jobKey);
    return storeProvider.getCronJobStore().fetchJob(jobKey);
  }

  private static class TaskValidationException extends Exception {
    TaskValidationException(String message) {
      super(message);
    }
  }

  private void validateTaskLimits(
      int totalInstances,
      QuotaCheckResult quotaCheck) throws TaskValidationException {

    if (totalInstances <= 0 || totalInstances > thresholds.getMaxTasksPerJob()) {
      throw new TaskValidationException(String.format(
          "Instance count must be between 1 and %d inclusive.", thresholds.getMaxTasksPerJob()));
    }

    if (quotaCheck.getResult() == INSUFFICIENT_QUOTA) {
      throw new TaskValidationException("Insufficient resource quota: "
          + quotaCheck.getDetails().orElse(""));
    }
  }

  private static Set<InstanceTaskConfig> buildInitialState(Map<Integer, ITaskConfig> tasks) {
    // Translate tasks into instance IDs.
    Multimap<ITaskConfig, Integer> instancesByConfig = HashMultimap.create();
    Multimaps.invertFrom(Multimaps.forMap(tasks), instancesByConfig);

    // Reduce instance IDs into contiguous ranges.
    Map<ITaskConfig, Set<Range<Integer>>> rangesByConfig =
        Maps.transformValues(instancesByConfig.asMap(), Numbers::toRanges);

    ImmutableSet.Builder<InstanceTaskConfig> builder = ImmutableSet.builder();
    for (Map.Entry<ITaskConfig, Set<Range<Integer>>> entry : rangesByConfig.entrySet()) {
      builder.add(new InstanceTaskConfig()
          .setTask(entry.getKey().newBuilder())
          .setInstances(IRange.toBuildersSet(convertRanges(entry.getValue()))));
    }

    return builder.build();
  }

  @Override
  public Response startJobUpdate(JobUpdateRequest mutableRequest, @Nullable String message) {
    requireNonNull(mutableRequest);

    if (!mutableRequest.getTaskConfig().isIsService()) {
      return invalidRequest(NON_SERVICE_TASK);
    }

    int totalInstancesFromGroups;
    JobUpdateSettings settings = requireNonNull(mutableRequest.getSettings());

    // Gracefully handle a client sending an update with an older thrift schema
    // TODO(rdelvalle): Remove after version 0.22.0 ships
    ThriftBackfill.backfillUpdateStrategy(settings);

    // Keep old job schema in case we want to revert to a lower version of Aurora that doesn't
    // support variable update group sizes
    if (settings.getUpdateStrategy().isSetQueueStrategy()) {
      totalInstancesFromGroups = settings.getUpdateStrategy().getQueueStrategy().getGroupSize();
    } else if (settings.getUpdateStrategy().isSetBatchStrategy()) {
      totalInstancesFromGroups = settings.getUpdateStrategy().getBatchStrategy().getGroupSize();
    } else if (settings.getUpdateStrategy().isSetVarBatchStrategy()) {
      VariableBatchJobUpdateStrategy strategy = settings.getUpdateStrategy().getVarBatchStrategy();

      if (strategy.getGroupSizes().stream().anyMatch(x -> x <= 0)) {
        return invalidRequest(INVALID_GROUP_SIZE);
      }

      totalInstancesFromGroups = strategy.getGroupSizes().stream().reduce(0, Integer::sum);
    } else {
      return invalidRequest(UNKNOWN_UPDATE_STRATEGY);
    }

    if (totalInstancesFromGroups <= 0) {
      return invalidRequest(NO_INSTANCES_MODIFIED);
    }

    if (settings.getMaxPerInstanceFailures() < 0) {
      return invalidRequest(INVALID_MAX_INSTANCE_FAILURES);
    }

    if (settings.getMaxFailedInstances() < 0) {
      return invalidRequest(INVALID_MAX_FAILED_INSTANCES);
    }

    if (settings.getMaxPerInstanceFailures() * mutableRequest.getInstanceCount()
            > thresholds.getMaxUpdateInstanceFailures()) {
      return invalidRequest(TOO_MANY_POTENTIAL_FAILED_INSTANCES);
    }

    if (settings.getMinWaitInInstanceRunningMs() < 0) {
      return invalidRequest(INVALID_MIN_WAIT_TO_RUNNING);
    }

    if (settings.getBlockIfNoPulsesAfterMs() < 0) {
      return invalidRequest(INVALID_PULSE_TIMEOUT);
    }

    if (settings.isSlaAware() && !mutableRequest.getTaskConfig().isSetSlaPolicy()) {
      return invalidRequest(INVALID_SLA_AWARE_UPDATE);
    }

    IJobUpdateRequest request;
    try {
      request = IJobUpdateRequest.build(
          new JobUpdateRequest(mutableRequest)
              .setTaskConfig(configurationManager.validateAndPopulate(
                  ITaskConfig.build(mutableRequest.getTaskConfig()),
                  mutableRequest.getInstanceCount())
              .newBuilder()));
    } catch (TaskDescriptionException e) {
      return error(INVALID_REQUEST, e);
    }

    return storage.write(storeProvider -> {
      IJobKey job = request.getTaskConfig().getJob();
      if (getCronJob(storeProvider, job).isPresent()) {
        return invalidRequest(NO_CRON);
      }

      String updateId = uuidGenerator.createNew().toString();
      IJobUpdateSettings settings1 = request.getSettings();

      JobDiff diff = JobDiff.compute(
          storeProvider.getTaskStore(),
          job,
          JobDiff.asMap(request.getTaskConfig(), request.getInstanceCount()),
          settings1.getUpdateOnlyTheseInstances());

      Set<Integer> invalidScope = diff.getOutOfScopeInstances(
          Numbers.rangesToInstanceIds(settings1.getUpdateOnlyTheseInstances()));
      if (!invalidScope.isEmpty()) {
        return invalidRequest(
            "The update request attempted to update specific instances,"
                + " but some are irrelevant to the update and current job state: "
                + invalidScope);
      }

      if (diff.isNoop()) {
        return addMessage(empty(), OK, NOOP_JOB_UPDATE_MESSAGE);
      }

      JobUpdateInstructions instructions = new JobUpdateInstructions()
          .setSettings(settings1.newBuilder())
          .setInitialState(buildInitialState(diff.getReplacedInstances()));

      Set<Integer> replacements = diff.getReplacementInstances();
      if (!replacements.isEmpty()) {
        instructions.setDesiredState(
            new InstanceTaskConfig()
                .setTask(request.getTaskConfig().newBuilder())
                .setInstances(IRange.toBuildersSet(convertRanges(toRanges(replacements)))));
      }

      String remoteUserName = auditMessages.getRemoteUserName();
      IJobUpdate update = IJobUpdate.build(new JobUpdate()
          .setSummary(new JobUpdateSummary()
              .setKey(new JobUpdateKey(job.newBuilder(), updateId))
              .setUser(remoteUserName)
              .setMetadata(IMetadata.toBuildersSet(request.getMetadata())))
          .setInstructions(instructions));

      Response response = empty();
      try {
        validateTaskLimits(
            request.getInstanceCount(),
            quotaManager.checkJobUpdate(update, storeProvider));

        jobUpdateController.start(
            update,
            new AuditData(remoteUserName, Optional.ofNullable(message)));
        startJobUpdateCounter.addAndGet(request.getInstanceCount());
        return response.setResponseCode(OK)
            .setResult(Result.startJobUpdateResult(
                new StartJobUpdateResult(update.getSummary().getKey().newBuilder())
                    .setUpdateSummary(update.getSummary().newBuilder())));
      } catch (UpdateInProgressException e) {
        return error(INVALID_REQUEST, e)
            .setResult(Result.startJobUpdateResult(
                new StartJobUpdateResult(e.getInProgressUpdateSummary().getKey().newBuilder())
                    .setUpdateSummary(e.getInProgressUpdateSummary().newBuilder())));
      } catch (UpdateStateException | TaskValidationException e) {
        return error(INVALID_REQUEST, e);
      }
    });
  }

  private Response changeJobUpdateState(
      JobUpdateKey mutableKey,
      JobUpdateStateChange change,
      Optional<String> message) {

    IJobUpdateKey key = IJobUpdateKey.build(mutableKey);
    JobKeys.assertValid(key.getJob());
    return storage.write(storeProvider -> {
      try {
        change.modifyUpdate(
            jobUpdateController,
            key,
            new AuditData(auditMessages.getRemoteUserName(), message));
        return ok();
      } catch (UpdateStateException e) {
        return error(INVALID_REQUEST, e);
      }
    });
  }

  private interface JobUpdateStateChange {
    void modifyUpdate(JobUpdateController controller, IJobUpdateKey key, AuditData auditData)
        throws UpdateStateException;
  }

  @Override
  public Response pauseJobUpdate(JobUpdateKey mutableKey, @Nullable String message) {
    return changeJobUpdateState(
        mutableKey,
        JobUpdateController::pause,
        Optional.ofNullable(message));
  }

  @Override
  public Response resumeJobUpdate(JobUpdateKey mutableKey, @Nullable String message) {
    return changeJobUpdateState(
        mutableKey,
        JobUpdateController::resume,
        Optional.ofNullable(message));
  }

  @Override
  public Response abortJobUpdate(JobUpdateKey mutableKey, @Nullable String message) {
    return changeJobUpdateState(
        mutableKey,
        JobUpdateController::abort,
        Optional.ofNullable(message));
  }

  @Override
  public Response rollbackJobUpdate(JobUpdateKey mutableKey, @Nullable String message) {
    return changeJobUpdateState(
        mutableKey,
        JobUpdateController::rollback,
        Optional.ofNullable(message));
  }

  @Override
  public Response pulseJobUpdate(JobUpdateKey mutableUpdateKey) {
    IJobUpdateKey updateKey = validateJobUpdateKey(mutableUpdateKey);
    try {
      JobUpdatePulseStatus result = jobUpdateController.pulse(updateKey);
      return ok(Result.pulseJobUpdateResult(new PulseJobUpdateResult(result)));
    } catch (UpdateStateException e) {
      return error(INVALID_REQUEST, e);
    }
  }

  @Override
  public Response pruneTasks(TaskQuery query) throws TException {
    if (query.isSetStatuses() && query.getStatuses().stream().anyMatch(ACTIVE_STATES::contains)) {
      return error("Tasks in non-terminal state cannot be pruned.");
    } else if (!query.isSetStatuses()) {
      query.setStatuses(TERMINAL_STATES);
    }

    Iterable<IScheduledTask> tasks = storage.read(storeProvider ->
        storeProvider.getTaskStore().fetchTasks(Query.arbitrary(query)));
    // For some reason fetchTasks ignores the offset/limit options of a TaskQuery. So we have to
    // manually apply the limit here. To be fixed in AURORA-1892.
    if (query.isSetLimit()) {
      tasks = Iterables.limit(tasks, query.getLimit());
    }

    Iterable<String> taskIds = Iterables.transform(
        tasks,
        task -> task.getAssignedTask().getTaskId());

    return storage.write(storeProvider -> {
      Set<String> taskIdsSet = Sets.newHashSet(taskIds);
      stateManager.deleteTasks(storeProvider, taskIdsSet);
      pruneTasksCounter.addAndGet(taskIdsSet.size());
      return ok();
    });
  }

  @ThriftWorkload
  @Override
  public Response getJobUpdateSummaries(JobUpdateQuery mutableQuery) throws TException {
    return readOnlyScheduler.getJobUpdateSummaries(mutableQuery);
  }

  @ThriftWorkload
  @Override
  public Response getJobUpdateDetails(JobUpdateKey key, JobUpdateQuery query) throws TException {
    return readOnlyScheduler.getJobUpdateDetails(key, query);
  }

  private static IJobUpdateKey validateJobUpdateKey(JobUpdateKey mutableKey) {
    IJobUpdateKey key = IJobUpdateKey.build(mutableKey);
    JobKeys.assertValid(key.getJob());
    checkNotBlank(key.getId());
    return key;
  }

  @VisibleForTesting
  static String noCronScheduleMessage(IJobKey jobKey) {
    return String.format("Job %s has no cron schedule", JobKeys.canonicalString(jobKey));
  }

  @VisibleForTesting
  static String notScheduledCronMessage(IJobKey jobKey) {
    return String.format("Job %s is not scheduled with cron", JobKeys.canonicalString(jobKey));
  }

  @VisibleForTesting
  static String jobAlreadyExistsMessage(IJobKey jobKey) {
    return String.format("Job %s already exists", JobKeys.canonicalString(jobKey));
  }

  @VisibleForTesting
  static final String NO_TASKS_TO_KILL_MESSAGE = "No tasks to kill.";

  @VisibleForTesting
  static final String NOOP_JOB_UPDATE_MESSAGE = "Job is unchanged by proposed update.";

  @VisibleForTesting
  static final String NO_CRON = "Cron jobs may only be created/updated by calling scheduleCronJob.";

  @VisibleForTesting
  static final String NO_INSTANCES_MODIFIED = "Update results in no instance being modified.";

  @VisibleForTesting
  static final String NON_SERVICE_TASK = "Updates are not supported for non-service tasks.";

  @VisibleForTesting
  static final String INVALID_GROUP_SIZE = "All update group sizes must be positive.";

  @VisibleForTesting
  static final String INVALID_MAX_FAILED_INSTANCES = "maxFailedInstances must be non-negative.";

  @VisibleForTesting
  static final String TOO_MANY_POTENTIAL_FAILED_INSTANCES = "Your update allows too many failures "
      + "to occur, consider decreasing the per-instance failures or maxFailedInstances.";

  @VisibleForTesting
  static final String INVALID_MAX_INSTANCE_FAILURES
      = "maxPerInstanceFailures must be non-negative.";

  @VisibleForTesting
  static final String INVALID_MIN_WAIT_TO_RUNNING =
      "minWaitInInstanceRunningMs must be non-negative.";

  @VisibleForTesting
  static final String INVALID_PULSE_TIMEOUT = "blockIfNoPulsesAfterMs must be positive.";

  @VisibleForTesting
  static final String INVALID_INSTANCE_ID = "No active task found for a given instance ID.";

  @VisibleForTesting
  static final String INVALID_INSTANCE_COUNT = "Instance count must be positive.";

  @VisibleForTesting
  static final String UNKNOWN_UPDATE_STRATEGY = "Update strategy provided is unknown.";

  @VisibleForTesting
  static final String INVALID_SLA_AWARE_UPDATE = "slaAware is true, but no task slaPolicy is "
      + "specified.";
}