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

import com.datastax.oss.driver.api.core.DriverTimeoutException;
import com.datastax.oss.driver.api.core.RequestThrottlingException;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metrics.DefaultSessionMetric;
import com.datastax.oss.driver.api.core.session.throttling.RequestThrottler;
import com.datastax.oss.driver.api.core.session.throttling.Throttled;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.internal.core.control.ControlConnection;
import com.datastax.oss.driver.internal.core.metrics.SessionMetricUpdater;
import com.datastax.oss.driver.internal.core.pool.ChannelPool;
import com.datastax.oss.driver.internal.core.session.DefaultSession;
import com.datastax.oss.protocol.internal.Message;
import com.datastax.oss.protocol.internal.response.Result;
import com.datastax.oss.protocol.internal.response.result.Prepared;
import com.datastax.oss.protocol.internal.response.result.Rows;
import edu.umd.cs.findbugs.annotations.NonNull;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import net.jcip.annotations.ThreadSafe;

@ThreadSafe
public class ThrottledAdminRequestHandler<ResultT> extends AdminRequestHandler<ResultT>
    implements Throttled {

  
Params:
shouldPreAcquireId – whether to call DriverChannel.preAcquireId() before sending the request. This must be false if you obtained the connection from a pool (ChannelPool.next(), or DefaultSession.getChannel(Node, String)). It must be
    true if you are using a standalone channel (e.g. in ControlConnection or one of its auxiliary components)./**
   * @param shouldPreAcquireId whether to call {@link DriverChannel#preAcquireId()} before sending
   *     the request. This <b>must be false</b> if you obtained the connection from a pool ({@link
   *     ChannelPool#next()}, or {@link DefaultSession#getChannel(Node, String)}). It <b>must be
   *     true</b> if you are using a standalone channel (e.g. in {@link ControlConnection} or one of
   *     its auxiliary components).
   */
  public static ThrottledAdminRequestHandler<AdminResult> query(
      DriverChannel channel,
      boolean shouldPreAcquireId,
      Message message,
      Map<String, ByteBuffer> customPayload,
      Duration timeout,
      RequestThrottler throttler,
      SessionMetricUpdater metricUpdater,
      String logPrefix,
      String debugString) {
    return new ThrottledAdminRequestHandler<>(
        channel,
        shouldPreAcquireId,
        message,
        customPayload,
        timeout,
        throttler,
        metricUpdater,
        logPrefix,
        debugString,
        Rows.class);
  }

  
Params:
shouldPreAcquireId – whether to call DriverChannel.preAcquireId() before sending the request. See query(DriverChannel, boolean, Message, Map<String,ByteBuffer>, Duration, RequestThrottler, SessionMetricUpdater, String, String) for more explanations./**
   * @param shouldPreAcquireId whether to call {@link DriverChannel#preAcquireId()} before sending
   *     the request. See {@link #query(DriverChannel, boolean, Message, Map, Duration,
   *     RequestThrottler, SessionMetricUpdater, String, String)} for more explanations.
   */
  public static ThrottledAdminRequestHandler<ByteBuffer> prepare(
      DriverChannel channel,
      boolean shouldPreAcquireId,
      Message message,
      Map<String, ByteBuffer> customPayload,
      Duration timeout,
      RequestThrottler throttler,
      SessionMetricUpdater metricUpdater,
      String logPrefix) {
    return new ThrottledAdminRequestHandler<>(
        channel,
        shouldPreAcquireId,
        message,
        customPayload,
        timeout,
        throttler,
        metricUpdater,
        logPrefix,
        message.toString(),
        Prepared.class);
  }

  private final long startTimeNanos;
  private final RequestThrottler throttler;
  private final SessionMetricUpdater metricUpdater;

  protected ThrottledAdminRequestHandler(
      DriverChannel channel,
      boolean preAcquireId,
      Message message,
      Map<String, ByteBuffer> customPayload,
      Duration timeout,
      RequestThrottler throttler,
      SessionMetricUpdater metricUpdater,
      String logPrefix,
      String debugString,
      Class<? extends Result> expectedResponseType) {
    super(
        channel,
        preAcquireId,
        message,
        customPayload,
        timeout,
        logPrefix,
        debugString,
        expectedResponseType);
    this.startTimeNanos = System.nanoTime();
    this.throttler = throttler;
    this.metricUpdater = metricUpdater;
  }

  @Override
  public CompletionStage<ResultT> start() {
    // Don't write request yet, wait for green light from throttler
    throttler.register(this);
    return result;
  }

  @Override
  public void onThrottleReady(boolean wasDelayed) {
    if (wasDelayed) {
      metricUpdater.updateTimer(
          DefaultSessionMetric.THROTTLING_DELAY,
          null,
          System.nanoTime() - startTimeNanos,
          TimeUnit.NANOSECONDS);
    }
    super.start();
  }

  @Override
  public void onThrottleFailure(@NonNull RequestThrottlingException error) {
    metricUpdater.incrementCounter(DefaultSessionMetric.THROTTLING_ERRORS, null);
    setFinalError(error);
  }

  @Override
  protected boolean setFinalResult(ResultT result) {
    boolean wasSet = super.setFinalResult(result);
    if (wasSet) {
      throttler.signalSuccess(this);
    }
    return wasSet;
  }

  @Override
  protected boolean setFinalError(Throwable error) {
    boolean wasSet = super.setFinalError(error);
    if (wasSet) {
      if (error instanceof DriverTimeoutException) {
        throttler.signalTimeout(this);
      } else if (!(error instanceof RequestThrottlingException)) {
        throttler.signalError(this, error);
      }
    }
    return wasSet;
  }
}