/*
 * 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 static org.assertj.core.api.Assertions.assertThat;

import com.datastax.driver.core.exceptions.DriverException;
import com.datastax.driver.core.policies.DelegatingLoadBalancingPolicy;
import com.datastax.driver.core.policies.DelegatingSpeculativeExecutionPolicy;
import com.datastax.driver.core.policies.NoSpeculativeExecutionPolicy;
import com.datastax.driver.core.policies.RetryPolicy;
import com.datastax.driver.core.policies.RoundRobinPolicy;
import com.datastax.driver.core.utils.CassandraVersion;
import java.util.Iterator;
import java.util.concurrent.atomic.AtomicInteger;
import org.testng.annotations.Test;

public class StatementWrapperTest extends CCMTestsSupport {

  private static final String INSERT_QUERY = "insert into test (k, v) values (?, ?)";
  private static final String SELECT_QUERY = "select * from test where k = ?";

  private static final String INSERT_MULTIPAGE_QUERY =
      "insert into test_multipage (k, v) values (?, ?)";
  private static final String SELECT_MULTIPAGE_QUERY = "select * from test_multipage where k = ?";

  CustomLoadBalancingPolicy loadBalancingPolicy = new CustomLoadBalancingPolicy();
  CustomSpeculativeExecutionPolicy speculativeExecutionPolicy =
      new CustomSpeculativeExecutionPolicy();
  CustomRetryPolicy retryPolicy = new CustomRetryPolicy();

  @Override
  public void onTestContextInitialized() {
    execute("create table test (k text primary key, v int)");
    execute("create table test_multipage (k text, v int, primary key (k, v))");
  }

  @Override
  public Cluster.Builder createClusterBuilder() {
    return super.createClusterBuilder()
        .withLoadBalancingPolicy(loadBalancingPolicy)
        .withSpeculativeExecutionPolicy(speculativeExecutionPolicy)
        .withRetryPolicy(retryPolicy);
  }

  @Test(groups = "short")
  public void should_pass_wrapped_statement_to_load_balancing_policy() {
    loadBalancingPolicy.customStatementsHandled.set(0);

    SimpleStatement s = new SimpleStatement("select * from system.local");
    session().execute(s);
    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(0);

    session().execute(new CustomStatement(s));
    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);
  }

  @Test(groups = "short")
  @CassandraVersion(value = "2.0.0")
  public void should_reuse_wrapped_simple_statement_for_multipage_query() {
    loadBalancingPolicy.customStatementsHandled.set(0);

    for (int v = 1; v <= 100; v++)
      session().execute(new SimpleStatement(INSERT_MULTIPAGE_QUERY, "key_simple_multipage", v));

    SimpleStatement s = new SimpleStatement(SELECT_MULTIPAGE_QUERY, "key_simple_multipage");
    s.setFetchSize(1);

    CustomStatement customStatement = new CustomStatement(s);

    ResultSet rs = session().execute(customStatement);
    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

    Iterator<Row> it = rs.iterator();

    assertThat(it.hasNext()).isTrue();
    it.next();
    assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

    assertThat(it.hasNext()).isTrue();
    it.next();
    assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);

    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(2);
  }

  @Test(groups = "short")
  @CassandraVersion(value = "2.0.0")
  public void should_reuse_wrapped_bound_statement_for_multipage_query() {
    loadBalancingPolicy.customStatementsHandled.set(0);

    for (int v = 1; v <= 100; v++)
      session().execute(new SimpleStatement(INSERT_MULTIPAGE_QUERY, "key_prepared_multipage", v));

    PreparedStatement ps = session().prepare(SELECT_MULTIPAGE_QUERY);
    BoundStatement bs = ps.bind("key_prepared_multipage");
    bs.setFetchSize(1);

    CustomStatement customStatement = new CustomStatement(bs);

    ResultSet rs = session().execute(customStatement);
    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

    Iterator<Row> it = rs.iterator();

    assertThat(it.hasNext()).isTrue();
    it.next();
    assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

    assertThat(it.hasNext()).isTrue();
    it.next();
    assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);

    assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(2);
  }

  @Test(groups = "short")
  public void should_pass_wrapped_statement_to_speculative_execution_policy() {
    speculativeExecutionPolicy.customStatementsHandled.set(0);

    SimpleStatement s = new SimpleStatement("select * from system.local");
    session().execute(s);
    assertThat(speculativeExecutionPolicy.customStatementsHandled.get()).isEqualTo(0);

    session().execute(new CustomStatement(s));
    assertThat(speculativeExecutionPolicy.customStatementsHandled.get()).isEqualTo(1);
  }

  @Test(groups = "short")
  public void should_pass_wrapped_statement_to_retry_policy() {
    retryPolicy.customStatementsHandled.set(0);

    // Set CL TWO with only one node, so the statement will always cause UNAVAILABLE,
    // which our custom policy ignores.
    Statement s =
        new SimpleStatement("select * from system.local").setConsistencyLevel(ConsistencyLevel.TWO);

    session().execute(s);
    assertThat(retryPolicy.customStatementsHandled.get()).isEqualTo(0);

    session().execute(new CustomStatement(s));
    assertThat(retryPolicy.customStatementsHandled.get()).isEqualTo(1);
  }

  @CassandraVersion("2.0.0")
  @Test(groups = "short")
  public void should_execute_wrapped_simple_statement() {
    session().execute(new CustomStatement(new SimpleStatement(INSERT_QUERY, "key_simple", 1)));

    ResultSet rs =
        session().execute(new CustomStatement(new SimpleStatement(SELECT_QUERY, "key_simple")));
    assertThat(rs.one().getInt("v")).isEqualTo(1);
  }

  @Test(groups = "short")
  public void should_execute_wrapped_bound_statement() {
    PreparedStatement preparedStatement = session().prepare(new SimpleStatement(INSERT_QUERY));
    session().execute(new CustomStatement(preparedStatement.bind("key_bound", 1)));

    preparedStatement = session().prepare(new SimpleStatement(SELECT_QUERY));
    ResultSet rs = session().execute(new CustomStatement(preparedStatement.bind("key_bound")));
    assertThat(rs.one().getInt("v")).isEqualTo(1);
  }

  @CassandraVersion("2.0.0")
  @Test(groups = "short")
  public void should_execute_wrapped_batch_statement() {
    BatchStatement batchStatement = new BatchStatement();
    batchStatement.add(new SimpleStatement(INSERT_QUERY, "key_batch", 1));

    session().execute(new CustomStatement(batchStatement));

    ResultSet rs = session().execute(SELECT_QUERY, "key_batch");
    assertThat(rs.one().getInt("v")).isEqualTo(1);
  }

  @CassandraVersion("2.0.0")
  @Test(groups = "short")
  public void should_add_wrapped_batch_statement_to_batch_statement() {
    BatchStatement batchStatementForWrapping = new BatchStatement();
    batchStatementForWrapping.add(new SimpleStatement(INSERT_QUERY, "key1", 1));

    BatchStatement batchStatement = new BatchStatement();
    batchStatement.add(new CustomStatement(new SimpleStatement(INSERT_QUERY, "key2", 2)));
    batchStatement.add(new CustomStatement(batchStatementForWrapping));

    session().execute(batchStatement);

    ResultSet rs = session().execute(SELECT_QUERY, "key1");
    assertThat(rs.one().getInt("v")).isEqualTo(1);

    rs = session().execute(SELECT_QUERY, "key2");
    assertThat(rs.one().getInt("v")).isEqualTo(2);
  }

  /** A custom wrapper that's just used to mark statements. */
  static class CustomStatement extends StatementWrapper {
    protected CustomStatement(Statement wrapped) {
      super(wrapped);
    }
  }

  /** A load balancing policy that counts how many times it has seen the custom wrapper */
  static class CustomLoadBalancingPolicy extends DelegatingLoadBalancingPolicy {
    final AtomicInteger customStatementsHandled = new AtomicInteger();

    public CustomLoadBalancingPolicy() {
      super(new RoundRobinPolicy());
    }

    @Override
    public Iterator<Host> newQueryPlan(String loggedKeyspace, Statement statement) {
      if (statement instanceof CustomStatement) customStatementsHandled.incrementAndGet();
      return super.newQueryPlan(loggedKeyspace, statement);
    }
  }

  /** A speculative execution policy that counts how many times it has seen the custom wrapper */
  static class CustomSpeculativeExecutionPolicy extends DelegatingSpeculativeExecutionPolicy {
    final AtomicInteger customStatementsHandled = new AtomicInteger();

    protected CustomSpeculativeExecutionPolicy() {
      super(NoSpeculativeExecutionPolicy.INSTANCE);
    }

    @Override
    public SpeculativeExecutionPlan newPlan(String loggedKeyspace, Statement statement) {
      if (statement instanceof CustomStatement) customStatementsHandled.incrementAndGet();
      return super.newPlan(loggedKeyspace, statement);
    }
  }

  /**
   * A retry policy that counts how many times it has seen the custom wrapper for UNAVAILABLE
   * errors.
   */
  static class CustomRetryPolicy implements RetryPolicy {
    final AtomicInteger customStatementsHandled = new AtomicInteger();

    @Override
    public RetryDecision onUnavailable(
        Statement statement,
        ConsistencyLevel cl,
        int requiredReplica,
        int aliveReplica,
        int nbRetry) {
      if (statement instanceof CustomStatement) customStatementsHandled.incrementAndGet();
      return RetryDecision.ignore();
    }

    @Override
    public RetryDecision onReadTimeout(
        Statement statement,
        ConsistencyLevel cl,
        int requiredResponses,
        int receivedResponses,
        boolean dataRetrieved,
        int nbRetry) {
      return RetryDecision.rethrow();
    }

    @Override
    public RetryDecision onWriteTimeout(
        Statement statement,
        ConsistencyLevel cl,
        WriteType writeType,
        int requiredAcks,
        int receivedAcks,
        int nbRetry) {
      return RetryDecision.rethrow();
    }

    @Override
    public RetryDecision onRequestError(
        Statement statement, ConsistencyLevel cl, DriverException e, int nbRetry) {
      return RetryDecision.tryNextHost(cl);
    }

    @Override
    public void init(Cluster cluster) {
      // nothing to do
    }

    @Override
    public void close() {
      // nothing to do
    }
  }
}
