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

import java.util.List;
import java.util.Objects;
import java.util.Optional;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;

import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Time;
import org.apache.aurora.common.util.testing.FakeClock;
import org.apache.aurora.gen.AssignedTask;
import org.apache.aurora.gen.Constraint;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.TaskConfig;
import org.apache.aurora.gen.TaskConstraint;
import org.apache.aurora.gen.TaskEvent;
import org.apache.aurora.gen.ValueConstraint;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.junit.Test;

import static org.apache.aurora.gen.Resource.numCpus;
import static org.apache.aurora.gen.ScheduleStatus.ASSIGNED;
import static org.apache.aurora.gen.ScheduleStatus.FAILED;
import static org.apache.aurora.gen.ScheduleStatus.FINISHED;
import static org.apache.aurora.gen.ScheduleStatus.KILLING;
import static org.apache.aurora.gen.ScheduleStatus.PENDING;
import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
import static org.apache.aurora.gen.ScheduleStatus.STARTING;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result.EVALUATE_AFTER_MIN_RUNNING_MS;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result.EVALUATE_ON_STATE_CHANGE;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result.KILL_TASK_AND_EVALUATE_ON_STATE_CHANGE;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result.KILL_TASK_WITH_RESERVATION_AND_EVALUATE_ON_STATE_CHANGE;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result.REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE;
import static org.apache.aurora.scheduler.updater.StateEvaluator.Result.SUCCEEDED;
import static org.junit.Assert.assertEquals;

public class InstanceUpdaterTest {
  private static final Optional<ITaskConfig> NO_CONFIG = Optional.empty();

  private static final ITaskConfig OLD = ITaskConfig.build(new TaskConfig()
          .setResources(ImmutableSet.of(numCpus(1.0))));
  private static final ITaskConfig NEW = ITaskConfig.build(new TaskConfig()
          .setProduction(true)
          .setResources(ImmutableSet.of(numCpus(1.0))));
  private static final ITaskConfig NEW_EXTRA_RESOURCES = ITaskConfig.build(new TaskConfig()
      .setResources(ImmutableSet.of(numCpus(2.0))));
  private static final ITaskConfig NEW_DIFFERENT_CONSTRAINTS = ITaskConfig.build(new TaskConfig()
      .setConstraints(ImmutableSet.of(new Constraint("different",
          TaskConstraint.value(new ValueConstraint(false, ImmutableSet.of("test")))))));

  private static final Amount<Long, Time> MIN_RUNNING_TIME = Amount.of(1L, Time.MINUTES);
  private static final Amount<Long, Time> A_LONG_TIME = Amount.of(1L, Time.DAYS);

  private static class TestFixture {
    private final FakeClock clock;
    private final InstanceUpdater updater;
    private final TaskUtil taskUtil;
    private Optional<IScheduledTask> task = Optional.empty();

    TestFixture(Optional<ITaskConfig> newConfig, int maxToleratedFailures) {
      this.clock = new FakeClock();
      this.updater = new InstanceUpdater(newConfig, maxToleratedFailures, MIN_RUNNING_TIME, clock);
      this.taskUtil = new TaskUtil(clock);
    }

    TestFixture(ITaskConfig newConfig, int maxToleratedFailures) {
      this(Optional.of(newConfig), maxToleratedFailures);
    }

    void setActualState(ITaskConfig config) {
      this.task = Optional.of(taskUtil.makeTask(config, PENDING));
    }

    void setActualStateAbsent() {
      this.task = Optional.empty();
    }

    private Result changeStatusAndEvaluate(ScheduleStatus status) {
      ScheduledTask builder = task.get().newBuilder();
      if (builder.getStatus() != status) {
        // Only add a task event if this is a state change.
        builder.addToTaskEvents(new TaskEvent().setTimestamp(clock.nowMillis()).setStatus(status));
      }
      builder.setStatus(status);

      task = Optional.of(IScheduledTask.build(builder));
      return updater.evaluate(task);
    }

    void evaluateCurrentState(Result expectedResult) {
      assertEquals(expectedResult, updater.evaluate(task));
    }

    void evaluate(
        Result expectedResult,
        ScheduleStatus status,
        ScheduleStatus... statuses) {

      assertEquals(expectedResult, changeStatusAndEvaluate(status));
      for (ScheduleStatus s : statuses) {
        assertEquals(expectedResult, changeStatusAndEvaluate(s));
      }
    }

    void advanceTime(Amount<Long, Time> time) {
      clock.advance(time);
    }
  }

  @Test
  public void testSuccessfulUpdate() {
    TestFixture f = new TestFixture(NEW, 1);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_WITH_RESERVATION_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.advanceTime(MIN_RUNNING_TIME);
    f.evaluateCurrentState(SUCCEEDED);
  }

  @Test
  public void testUpdateWithResourceChange() {
    TestFixture f = new TestFixture(NEW_EXTRA_RESOURCES, 1);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.setActualState(NEW_EXTRA_RESOURCES);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.advanceTime(MIN_RUNNING_TIME);
    f.evaluateCurrentState(SUCCEEDED);
  }

  @Test
  public void testUpdateWithConstraintChange() {
    TestFixture f = new TestFixture(NEW_DIFFERENT_CONSTRAINTS, 1);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.setActualState(NEW_DIFFERENT_CONSTRAINTS);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.advanceTime(MIN_RUNNING_TIME);
    f.evaluateCurrentState(SUCCEEDED);
  }

  @Test
  public void testUpdateRetryOnTaskExit() {
    TestFixture f = new TestFixture(NEW, 1);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_WITH_RESERVATION_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, FAILED);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.advanceTime(MIN_RUNNING_TIME);
    f.evaluateCurrentState(SUCCEEDED);
  }

  @Test
  public void testUpdateRetryFailure() {
    TestFixture f = new TestFixture(NEW, 0);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_WITH_RESERVATION_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.evaluate(Result.FAILED_TERMINATED, FAILED);
  }

  @Test
  public void testNoopUpdate() {
    TestFixture f = new TestFixture(NEW, 1);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.advanceTime(MIN_RUNNING_TIME);
    f.evaluate(SUCCEEDED, RUNNING);
  }

  @Test
  public void testPointlessUpdate() {
    TestFixture f = new TestFixture(NO_CONFIG, 1);
    f.setActualStateAbsent();
    f.evaluateCurrentState(SUCCEEDED);
  }

  @Test
  public void testNoOldConfig() {
    TestFixture f = new TestFixture(NEW, 1);
    f.setActualStateAbsent();
    f.evaluateCurrentState(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.advanceTime(MIN_RUNNING_TIME);
    f.evaluateCurrentState(SUCCEEDED);
  }

  @Test
  public void testNoNewConfig() {
    TestFixture f = new TestFixture(NO_CONFIG, 1);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(SUCCEEDED, FINISHED);
    f.evaluateCurrentState(SUCCEEDED);
  }

  @Test
  public void testStuckInPending() {
    TestFixture f = new TestFixture(NEW, 1);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_WITH_RESERVATION_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING);
    f.advanceTime(A_LONG_TIME);
    f.evaluateCurrentState(EVALUATE_ON_STATE_CHANGE);
  }

  @Test
  public void testSlowToKill() {
    TestFixture f = new TestFixture(NEW, 1);
    f.setActualState(OLD);
    f.evaluate(KILL_TASK_WITH_RESERVATION_AND_EVALUATE_ON_STATE_CHANGE, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.evaluate(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING);
    f.advanceTime(A_LONG_TIME);
    f.evaluateCurrentState(EVALUATE_ON_STATE_CHANGE);
    f.setActualStateAbsent();
    f.evaluateCurrentState(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, KILLING);
    f.advanceTime(A_LONG_TIME);
    f.evaluateCurrentState(EVALUATE_ON_STATE_CHANGE);
  }

  @Test(expected = IllegalArgumentException.class)
  public void testInvalidInput() {
    TestFixture f = new TestFixture(NEW, 1);
    ScheduledTask noEvents = new TaskUtil(new FakeClock())
        .makeTask(OLD, RUNNING).newBuilder().setTaskEvents(ImmutableList.of());
    f.updater.evaluate(Optional.of(IScheduledTask.build(noEvents)));
  }

  @Test
  public void testOldTaskDies() {
    // If the original task dies, the updater should not add a replacement if the task will be
    // resuscitated.  Only a task that has passed through KILLING will not be resuscitated.

    TestFixture f = new TestFixture(NEW, 1);
    f.setActualState(OLD);
    // Task did not pass through KILLING, therefore will be rescheduled.
    f.evaluate(EVALUATE_ON_STATE_CHANGE, FINISHED);
  }

  @Test
  public void testOldTaskDiesAndRescheduled() {
    // Identical to testOldTaskDies, with the follow-through of rescheduling and updating.
    // If the original task dies, the updater should not add a replacement if the task will be
    // resuscitated.  Only a task that has passed through KILLING will not be resuscitated.

    TestFixture f = new TestFixture(NEW, 1);
    f.setActualState(OLD);

    // Task did not pass through KILLING, therefore will be rescheduled.
    f.evaluate(EVALUATE_ON_STATE_CHANGE, FINISHED);
    f.evaluate(KILL_TASK_WITH_RESERVATION_AND_EVALUATE_ON_STATE_CHANGE, PENDING);
    f.setActualStateAbsent();
    f.evaluateCurrentState(REPLACE_TASK_AND_EVALUATE_ON_STATE_CHANGE);
    f.setActualState(NEW);
    f.evaluate(EVALUATE_ON_STATE_CHANGE, PENDING, ASSIGNED, STARTING);
    f.evaluate(EVALUATE_AFTER_MIN_RUNNING_MS, RUNNING);
    f.advanceTime(MIN_RUNNING_TIME);
    f.evaluateCurrentState(SUCCEEDED);
  }

  static final class TaskUtil {
    private final FakeClock clock;

    TaskUtil(FakeClock clock) {
      this.clock = Objects.requireNonNull(clock);
    }

    IScheduledTask makeTask(ITaskConfig config, ScheduleStatus status) {
      List<TaskEvent> events = Lists.newArrayList();
      if (status != PENDING) {
        events.add(new TaskEvent().setTimestamp(clock.nowMillis()).setStatus(PENDING));
      }
      if (Tasks.isTerminated(status) || status == KILLING) {
        events.add(new TaskEvent().setTimestamp(clock.nowMillis()).setStatus(ASSIGNED));
        events.add(new TaskEvent().setTimestamp(clock.nowMillis()).setStatus(RUNNING));
      }

      events.add(new TaskEvent().setTimestamp(clock.nowMillis()).setStatus(status));

      return IScheduledTask.build(
          new ScheduledTask()
              .setStatus(status)
              .setTaskEvents(ImmutableList.copyOf(events))
              .setAssignedTask(
                  new AssignedTask()
                      .setTask(config.newBuilder())));
    }
  }
}