/*
 * Copyright DataStax, Inc.
 *
 * 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 com.datastax.oss.driver.internal.core.util.concurrent;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import com.datastax.oss.driver.shaded.guava.common.util.concurrent.Uninterruptibles;
import edu.umd.cs.findbugs.annotations.NonNull;
import io.netty.channel.DefaultEventLoop;
import io.netty.channel.EventLoopGroup;
import io.netty.util.concurrent.ScheduledFuture;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.FutureTask;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * Extend Netty's default event loop to capture scheduled tasks instead of running them. The tasks
 * can be checked later, and run manually.
 *
 * <p>Tasks submitted with {@link #execute(Runnable)} or {@link #submit(Callable)} are still
 * executed normally.
 *
 * <p>This is used to make unit tests independent of time.
 */
@SuppressWarnings("FunctionalInterfaceClash") // does not matter for test code
public class ScheduledTaskCapturingEventLoop extends DefaultEventLoop {

  private final BlockingQueue<CapturedTask<?>> capturedTasks = new ArrayBlockingQueue<>(100);

  public ScheduledTaskCapturingEventLoop(EventLoopGroup parent) {
    super(parent);
  }

  @NonNull
  @Override
  public <V> ScheduledFuture<V> schedule(Callable<V> callable, long delay, TimeUnit unit) {
    CapturedTask<V> task = new CapturedTask<>(callable, delay, unit);
    boolean added = capturedTasks.offer(task);
    assertThat(added).isTrue();
    return task.scheduledFuture;
  }

  @NonNull
  @Override
  public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit) {
    return schedule(
        () -> {
          command.run();
          return null;
        },
        delay,
        unit);
  }

  @NonNull
  @Override
  public ScheduledFuture<?> scheduleAtFixedRate(
      Runnable command, long initialDelay, long period, TimeUnit unit) {
    CapturedTask<?> task =
        new CapturedTask<>(
            () -> {
              command.run();
              return null;
            },
            initialDelay,
            period,
            unit);
    boolean added = capturedTasks.offer(task);
    assertThat(added).isTrue();
    return task.scheduledFuture;
  }

  @NonNull
  @Override
  public ScheduledFuture<?> scheduleWithFixedDelay(
      Runnable command, long initialDelay, long delay, TimeUnit unit) {
    throw new UnsupportedOperationException("Not supported yet");
  }

  public CapturedTask<?> nextTask() {
    try {
      return capturedTasks.poll(100, TimeUnit.MILLISECONDS);
    } catch (InterruptedException e) {
      fail("Unexpected interruption", e);
      throw new AssertionError();
    }
  }

  /**
   * Wait for any pending non-scheduled task (submitted with {@code submit}, {@code execute}, etc.)
   * to complete.
   */
  public void waitForNonScheduledTasks() {
    ScheduledFuture<Object> f = super.schedule(() -> null, 5, TimeUnit.NANOSECONDS);
    try {
      Uninterruptibles.getUninterruptibly(f, 1, TimeUnit.SECONDS);
    } catch (ExecutionException e) {
      fail("unexpected error", e.getCause());
    } catch (TimeoutException e) {
      fail("timed out while waiting for admin tasks to complete", e);
    }
  }

  public class CapturedTask<V> {
    private final FutureTask<V> futureTask;
    private final long initialDelay;
    private final long period;
    private final TimeUnit unit;

    @SuppressWarnings("unchecked")
    private final ScheduledFuture<V> scheduledFuture = mock(ScheduledFuture.class);

    CapturedTask(Callable<V> task, long initialDelay, TimeUnit unit) {
      this(task, initialDelay, -1, unit);
    }

    CapturedTask(Callable<V> task, long initialDelay, long period, TimeUnit unit) {
      this.futureTask = new FutureTask<>(task);
      this.initialDelay = initialDelay;
      this.period = period;
      this.unit = unit;

      // If the code under test cancels the scheduled future, cancel our task
      when(scheduledFuture.cancel(anyBoolean()))
          .thenAnswer(invocation -> futureTask.cancel(invocation.getArgument(0)));

      // Delegate methods of the scheduled future to our task (to be extended to more methods if
      // needed)
      when(scheduledFuture.isDone()).thenAnswer(invocation -> futureTask.isDone());
      when(scheduledFuture.isCancelled()).thenAnswer(invocation -> futureTask.isCancelled());
    }

    public void run() {
      submit(futureTask);
      waitForNonScheduledTasks();
    }

    public boolean isCancelled() {
      // futureTask.isCancelled() can create timing issues in CI environments, so give the
      // cancellation a short time to complete instead:
      try {
        futureTask.get(3, TimeUnit.SECONDS);
      } catch (CancellationException e) {
        return true;
      } catch (Exception e) {
        // ignore
      }
      return false;
    }

    public long getInitialDelay(TimeUnit targetUnit) {
      return targetUnit.convert(initialDelay, unit);
    }

    /** By convention, non-recurring tasks have a negative period */
    public long getPeriod(TimeUnit targetUnit) {
      return targetUnit.convert(period, unit);
    }
  }
}
