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

import java.util.Set;
import java.util.concurrent.ScheduledExecutorService;

import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.RateLimiter;

import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Time;
import org.apache.aurora.common.testing.easymock.EasyMockTest;
import org.apache.aurora.common.util.BackoffStrategy;
import org.apache.aurora.gen.AssignedTask;
import org.apache.aurora.gen.JobKey;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.TaskConfig;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.events.PubsubEvent;
import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
import org.apache.aurora.scheduler.scheduling.TaskGroups.TaskGroupBatchWorker;
import org.apache.aurora.scheduler.scheduling.TaskGroups.TaskGroupsSettings;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
import org.apache.aurora.scheduler.testing.FakeScheduledExecutor;
import org.apache.aurora.scheduler.testing.FakeStatsProvider;
import org.junit.Before;
import org.junit.Test;

import static org.apache.aurora.gen.ScheduleStatus.ASSIGNED;
import static org.apache.aurora.gen.ScheduleStatus.INIT;
import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
import static org.easymock.EasyMock.anyObject;
import static org.easymock.EasyMock.eq;
import static org.easymock.EasyMock.expect;
import static org.junit.Assert.assertEquals;

public class TaskGroupsTest extends EasyMockTest {
  private static final Amount<Long, Time> FIRST_SCHEDULE_DELAY = Amount.of(1L, Time.MILLISECONDS);
  private static final Amount<Long, Time> RESCHEDULE_DELAY = FIRST_SCHEDULE_DELAY;
  private static final IJobKey JOB_A = IJobKey.build(new JobKey("role", "test", "jobA"));
  private static final String TASK_A_ID = "a";
  private static final Set<String> SCHEDULED_RESULT = ImmutableSet.of(TASK_A_ID);

  private BackoffStrategy backoffStrategy;
  private TaskScheduler taskScheduler;
  private RateLimiter rateLimiter;
  private FakeScheduledExecutor clock;
  private RescheduleCalculator rescheduleCalculator;
  private TaskGroups taskGroups;
  private TaskGroupBatchWorker batchWorker;
  private StorageTestUtil storageUtil;
  private FakeStatsProvider statsProvider;

  @Before
  public void setUp() throws Exception {
    storageUtil = new StorageTestUtil(this);
    storageUtil.expectOperations();
    ScheduledExecutorService executor = createMock(ScheduledExecutorService.class);
    clock = FakeScheduledExecutor.fromScheduledExecutorService(executor);
    backoffStrategy = createMock(BackoffStrategy.class);
    taskScheduler = createMock(TaskScheduler.class);
    rateLimiter = createMock(RateLimiter.class);
    rescheduleCalculator = createMock(RescheduleCalculator.class);
    batchWorker = createMock(TaskGroupBatchWorker.class);
    statsProvider = new FakeStatsProvider();
    taskGroups = new TaskGroups(
        executor,
        new TaskGroupsSettings(FIRST_SCHEDULE_DELAY, backoffStrategy, rateLimiter, 2),
        taskScheduler,
        rescheduleCalculator,
        batchWorker,
        statsProvider);
  }

  @Test
  public void testEvaluatedAfterFirstSchedulePenalty() throws Exception {
    expect(rateLimiter.acquire()).andReturn(0D);
    expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of(TASK_A_ID))))
        .andReturn(SCHEDULED_RESULT);
    expectBatchExecute(batchWorker, storageUtil.storage, control, SCHEDULED_RESULT)
        .anyTimes();

    control.replay();

    taskGroups.taskChangedState(TaskStateChange.transition(makeTask(TASK_A_ID), INIT));
    clock.advance(FIRST_SCHEDULE_DELAY);
    assertEquals(0L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
  }

  @Test
  public void testTaskDeletedBeforeEvaluating() throws Exception {
    final IScheduledTask task = makeTask(TASK_A_ID);
    expect(rateLimiter.acquire()).andReturn(0.5D);
    expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of(TASK_A_ID))))
        .andAnswer(() -> {
          // Test a corner case where a task is deleted while it is being evaluated by the task
          // scheduler.  If not handled carefully, this could result in the scheduler trying again
          // later to satisfy the deleted task.
          taskGroups.tasksDeleted(new PubsubEvent.TasksDeleted(ImmutableSet.of(task)));

          return ImmutableSet.of();
        });
    expectBatchExecute(batchWorker, storageUtil.storage, control, ImmutableSet.of())
        .anyTimes();
    expect(backoffStrategy.calculateBackoffMs(FIRST_SCHEDULE_DELAY.as(Time.MILLISECONDS)))
        .andReturn(0L);

    control.replay();

    taskGroups.taskChangedState(TaskStateChange.transition(makeTask(Tasks.id(task)), INIT));
    clock.advance(FIRST_SCHEDULE_DELAY);
    assertEquals(1L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
  }

  @Test
  public void testEvaluatedOnStartup() throws Exception {
    expect(rateLimiter.acquire()).andReturn(0.000000001D);
    expect(rescheduleCalculator.getStartupScheduleDelayMs(makeTask(TASK_A_ID))).andReturn(1L);
    expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of(TASK_A_ID))))
        .andReturn(ImmutableSet.of(TASK_A_ID));
    expectBatchExecute(batchWorker, storageUtil.storage, control, SCHEDULED_RESULT)
        .anyTimes();

    control.replay();

    taskGroups.taskChangedState(TaskStateChange.initialized(makeTask(TASK_A_ID)));
    clock.advance(FIRST_SCHEDULE_DELAY);
    clock.advance(RESCHEDULE_DELAY);
    assertEquals(1L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
  }

  @Test
  public void testMultipleTasksAndResistStarvation() throws Exception {
    expect(rateLimiter.acquire()).andReturn(0.001D).times(2);
    expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of("a0", "a1"))))
        .andReturn(ImmutableSet.of("a0", "a1"));
    expect(taskScheduler.schedule(anyObject(), eq(ImmutableSet.of("b0"))))
        .andReturn(ImmutableSet.of("b0"));
    expectBatchExecute(
        batchWorker,
        storageUtil.storage,
        control,
        ImmutableSet.of("a0", "a1")).anyTimes();
    expectBatchExecute(batchWorker, storageUtil.storage, control, ImmutableSet.of("b0"))
        .anyTimes();

    control.replay();

    taskGroups.taskChangedState(TaskStateChange.transition(makeTask(JOB_A, "a0", 0), INIT));
    taskGroups.taskChangedState(TaskStateChange.transition(makeTask(JOB_A, "a1", 1), INIT));
    taskGroups.taskChangedState(TaskStateChange.transition(makeTask(JOB_A, "a2", 2), INIT));
    taskGroups.taskChangedState(TaskStateChange.transition(
        makeTask(IJobKey.build(JOB_A.newBuilder().setName("jobB")), "b0", 0), INIT));

    clock.advance(FIRST_SCHEDULE_DELAY);
    assertEquals(2L, statsProvider.getLongValue(TaskGroups.SCHEDULE_ATTEMPTS_BLOCKS));
  }

  @Test
  public void testNonPendingIgnored() {
    control.replay();

    IScheduledTask task =
        IScheduledTask.build(makeTask(TASK_A_ID).newBuilder().setStatus(ASSIGNED));
    taskGroups.taskChangedState(TaskStateChange.initialized(task));
  }

  private static IScheduledTask makeTask(String id) {
    return makeTask(JOB_A, id, 0);
  }

  private static IScheduledTask makeTask(IJobKey jobKey, String id, int instanceId) {
    return IScheduledTask.build(new ScheduledTask()
        .setStatus(ScheduleStatus.PENDING)
        .setAssignedTask(new AssignedTask()
            .setInstanceId(instanceId)
            .setTaskId(id)
            .setTask(new TaskConfig()
                .setJob(jobKey.newBuilder()))));
  }
}