/*
 * 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 io.netty.util.Timeout;
import io.netty.util.Timer;
import io.netty.util.TimerTask;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * Implementation of Netty's {@link io.netty.util.Timer Timer} interface to capture scheduled {@link
 * io.netty.util.Timeout Timeouts} instead of running them, so they can be run manually in tests.
 */
public class CapturingTimer implements Timer {

  private final ArrayBlockingQueue<CapturedTimeout> timeoutQueue = new ArrayBlockingQueue<>(16);

  @Override
  public Timeout newTimeout(TimerTask task, long delay, TimeUnit unit) {
    // delay and unit are not needed as the Timeout's TimerTask will be run manually
    CapturedTimeout timeout = new CapturedTimeout(task, this, delay, unit);
    // add the timeout to the queue
    timeoutQueue.add(timeout);
    return timeout;
  }

  /**
   * Retrieves the next scheduled Timeout. In tests, this will usually be a request timeout or a
   * speculative execution. Tests will need be able to predict the ordering as it is not easy to
   * tell from the returned Timeout itself.
   */
  public CapturedTimeout getNextTimeout() {
    return timeoutQueue.poll();
  }

  @Override
  public Set<Timeout> stop() {
    if (timeoutQueue.isEmpty()) {
      return Collections.emptySet();
    }
    Set<Timeout> timeoutsRemaining = new HashSet<>(timeoutQueue.size());
    for (Timeout t : timeoutQueue) {
      if (t != null) {
        t.cancel();
        timeoutsRemaining.add(t);
      }
    }
    return timeoutsRemaining;
  }

  /**
   * Implementation of Netty's {@link io.netty.util.Timeout Timeout} interface. It is just a simple
   * class that keeps track of the {@link io.netty.util.TimerTask TimerTask} and the {@link
   * io.netty.util.Timer Timer} implementation that should only be used in tests. The intended use
   * is to call the {@link io.netty.util.TimerTask#run(io.netty.util.Timeout) run()} method on the
   * TimerTask when you want to execute the task (so you don't have to depend on a real timer).
   *
   * <p>Example:
   *
   * <pre>{@code
   * // get the next timeout from the timer
   * Timeout t = timer.getNextTimeout();
   * // run the TimerTask associated with the timeout
   * t.task.run(t);
   * }</pre>
   */
  public static class CapturedTimeout implements Timeout {

    private final TimerTask task;
    private final CapturingTimer timer;
    private final long delay;
    private final TimeUnit unit;
    private final AtomicBoolean cancelled = new AtomicBoolean(false);

    private CapturedTimeout(TimerTask task, CapturingTimer timer, long delay, TimeUnit unit) {
      this.task = task;
      this.timer = timer;
      this.delay = delay;
      this.unit = unit;
    }

    @Override
    public Timer timer() {
      return timer;
    }

    @Override
    public TimerTask task() {
      return task;
    }

    public long getDelay(TimeUnit targetUnit) {
      return targetUnit.convert(delay, unit);
    }

    @Override
    public boolean isExpired() {
      return false;
    }

    @Override
    public boolean isCancelled() {
      return cancelled.get();
    }

    @Override
    public boolean cancel() {
      return cancelled.compareAndSet(false, true);
    }
  }
}
