/*
 * 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 com.google.common.util.concurrent.Uninterruptibles;
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.ExecutionException;
import java.util.concurrent.FutureTask;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.mockito.Mockito;

/**
 * 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);
  }

  @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;
  }

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

  @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;
  }

  @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 = Mockito.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
      Mockito.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)
      Mockito.when(scheduledFuture.isDone()).thenAnswer(invocation -> futureTask.isDone());
      Mockito.when(scheduledFuture.isCancelled())
          .thenAnswer(invocation -> futureTask.isCancelled());
    }

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

    public boolean isCancelled() {
      return futureTask.isCancelled();
    }

    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);
    }
  }
}

@@ -1131,16 +1131,67 @@
 llable;%0A
+import java.util.concurrent.CancellationException;%0A
 import j
@@ -5336,39 +5336,350 @@
     
-return futureTask.isCancelled()
+// futureTask.isCancelled() can create timing issues in CI environments, so give the%0A      // cancellation a short time to complete instead:%0A      try %7B%0A        futureTask.get(500, TimeUnit.MILLISECONDS);%0A      %7D catch (CancellationException e) %7B%0A        return true;%0A      %7D catch (Exception e) %7B%0A        // ignore%0A      %7D%0A      return false
 ;%0A  

Fix test failing in CI