/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.driver.core;

import com.datastax.driver.core.exceptions.AuthenticationException;
import com.datastax.driver.core.exceptions.ConnectionException;
import com.datastax.driver.core.exceptions.UnsupportedProtocolVersionException;
import com.datastax.driver.core.policies.ReconnectionPolicy;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.AbstractFuture;
import com.google.common.util.concurrent.ListenableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Manages periodic reconnection attempts after a host has been marked down.
 *
 * <p>Concurrent attempts are handled via the {@link #currentAttempt} reference passed to the
 * constructor. For a given reference, only one handler will run at a given time. Additional
 * handlers will cancel themselves if they find a previous handler running.
 *
 * <p>This class is designed for concurrency, but instances must not be shared: each thread creates
 * and starts its own private handler, all interactions happen through {@link #currentAttempt}.
 */
abstract class AbstractReconnectionHandler implements Runnable {

  private static final Logger logger = LoggerFactory.getLogger(AbstractReconnectionHandler.class);

  private final String name;
  private final ScheduledExecutorService executor;
  private final ReconnectionPolicy.ReconnectionSchedule schedule;
  /**
   * The future that is exposed to clients, representing completion of the current active handler
   */
  private final AtomicReference<ListenableFuture<?>> currentAttempt;

  @VisibleForTesting final HandlerFuture handlerFuture = new HandlerFuture();

  private final long initialDelayMs;

  private final CountDownLatch ready = new CountDownLatch(1);

  public AbstractReconnectionHandler(
      String name,
      ScheduledExecutorService executor,
      ReconnectionPolicy.ReconnectionSchedule schedule,
      AtomicReference<ListenableFuture<?>> currentAttempt) {
    this(name, executor, schedule, currentAttempt, -1);
  }

  public AbstractReconnectionHandler(
      String name,
      ScheduledExecutorService executor,
      ReconnectionPolicy.ReconnectionSchedule schedule,
      AtomicReference<ListenableFuture<?>> currentAttempt,
      long initialDelayMs) {
    this.name = name;
    this.executor = executor;
    this.schedule = schedule;
    this.currentAttempt = currentAttempt;
    this.initialDelayMs = initialDelayMs;
  }

  protected abstract Connection tryReconnect()
      throws ConnectionException, InterruptedException, UnsupportedProtocolVersionException,
          ClusterNameMismatchException;

  protected abstract void onReconnection(Connection connection);

  protected boolean onConnectionException(ConnectionException e, long nextDelayMs) {
    return true;
  }

  protected boolean onUnknownException(Exception e, long nextDelayMs) {
    return true;
  }

  // Retrying on authentication errors makes sense for applications that can update the credentials
  // at runtime, we don't want to force them
  // to restart.
  protected boolean onAuthenticationException(AuthenticationException e, long nextDelayMs) {
    return true;
  }

  // Retrying on these errors is unlikely to work
  protected boolean onUnsupportedProtocolVersionException(
      UnsupportedProtocolVersionException e, long nextDelayMs) {
    return false;
  }

  protected boolean onClusterNameMismatchException(
      ClusterNameMismatchException e, long nextDelayMs) {
    return false;
  }

  public void start() {
    long firstDelay = (initialDelayMs >= 0) ? initialDelayMs : schedule.nextDelayMs();
    logger.debug("First reconnection scheduled in {}ms", firstDelay);
    try {
      handlerFuture.nextTry = executor.schedule(this, firstDelay, TimeUnit.MILLISECONDS);

      while (true) {
        ListenableFuture<?> previous = currentAttempt.get();
        if (previous != null && !previous.isCancelled()) {
          logger.debug("Found another already active handler, cancelling");
          handlerFuture.cancel(false);
          break;
        }
        if (currentAttempt.compareAndSet(previous, handlerFuture)) {
          Host.statesLogger.debug("[{}] starting reconnection attempt", name);
          break;
        }
      }
      ready.countDown();
    } catch (RejectedExecutionException e) {
      // The executor has been shutdown, fair enough, just ignore
      logger.debug("Aborting reconnection handling since the cluster is shutting down");
    }
  }

  @Override
  public void run() {
    // Just make sure we don't start the first try too fast, in case we find out in start() that we
    // need to cancel ourselves
    try {
      ready.await();
    } catch (InterruptedException e) {
      // This can happen at shutdown
      Thread.currentThread().interrupt();
      return;
    }

    if (handlerFuture.isCancelled()) {
      logger.debug("Got cancelled, stopping");
      return;
    }

    try {
      onReconnection(tryReconnect());
      handlerFuture.markAsDone();
      currentAttempt.compareAndSet(handlerFuture, null);
      logger.debug("Reconnection successful, cleared the future");
    } catch (ConnectionException e) {
      long nextDelay = schedule.nextDelayMs();
      if (onConnectionException(e, nextDelay)) reschedule(nextDelay);
      else currentAttempt.compareAndSet(handlerFuture, null);
    } catch (AuthenticationException e) {
      logger.error(e.getMessage());
      long nextDelay = schedule.nextDelayMs();
      if (onAuthenticationException(e, nextDelay)) {
        reschedule(nextDelay);
      } else {
        logger.error(
            "Retries against {} have been suspended. It won't be retried unless the node is restarted.",
            e.getEndPoint());
        currentAttempt.compareAndSet(handlerFuture, null);
      }
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    } catch (UnsupportedProtocolVersionException e) {
      logger.error(e.getMessage());
      long nextDelay = schedule.nextDelayMs();
      if (onUnsupportedProtocolVersionException(e, nextDelay)) {
        reschedule(nextDelay);
      } else {
        logger.error(
            "Retries against {} have been suspended. It won't be retried unless the node is restarted.",
            e.getEndPoint());
        currentAttempt.compareAndSet(handlerFuture, null);
      }
    } catch (ClusterNameMismatchException e) {
      logger.error(e.getMessage());
      long nextDelay = schedule.nextDelayMs();
      if (onClusterNameMismatchException(e, nextDelay)) {
        reschedule(nextDelay);
      } else {
        logger.error(
            "Retries against {} have been suspended. It won't be retried unless the node is restarted.",
            e.endPoint);
        currentAttempt.compareAndSet(handlerFuture, null);
      }
    } catch (Exception e) {
      long nextDelay = schedule.nextDelayMs();
      if (onUnknownException(e, nextDelay)) reschedule(nextDelay);
      else currentAttempt.compareAndSet(handlerFuture, null);
    }
  }

  private void reschedule(long nextDelay) {
    // If we got cancelled during the failed reconnection attempt that lead here, don't reschedule
    if (handlerFuture.isCancelled()) {
      currentAttempt.compareAndSet(handlerFuture, null);
      return;
    }

    Host.statesLogger.debug("[{}] next reconnection attempt in {} ms", name, nextDelay);
    handlerFuture.nextTry = executor.schedule(this, nextDelay, TimeUnit.MILLISECONDS);
  }

  // The future that the handler exposes to its clients via currentAttempt
  @VisibleForTesting
  static class HandlerFuture extends AbstractFuture<Void> {
    // A future representing completion of the next task submitted to the executor
    volatile ScheduledFuture<?> nextTry;

    @Override
    public boolean cancel(boolean mayInterruptIfRunning) {
      // This is a check-then-act, so we may race with the scheduling of the first try, but in that
      // case
      // we'll re-check for cancellation when this first try starts running
      if (nextTry != null) {
        nextTry.cancel(mayInterruptIfRunning);
      }

      return super.cancel(mayInterruptIfRunning);
    }

    void markAsDone() {
      super.set(null);
    }
  }
}
