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

import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.isBootstrapping;
import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.noRows;
import static com.datastax.oss.simulacron.common.stubbing.PrimeDsl.when;
import static org.assertj.core.api.Assertions.assertThat;

import com.datastax.oss.driver.api.core.AllNodesFailedException;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverConfig;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.config.ProgrammaticDriverConfigLoaderBuilder;
import com.datastax.oss.driver.api.core.context.DriverContext;
import com.datastax.oss.driver.api.core.cql.ResultSet;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy;
import com.datastax.oss.driver.api.testinfra.loadbalancing.SortingLoadBalancingPolicy;
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
import com.datastax.oss.driver.api.testinfra.simulacron.QueryCounter;
import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule;
import com.datastax.oss.driver.categories.ParallelizableTests;
import com.datastax.oss.driver.internal.core.specex.ConstantSpeculativeExecutionPolicy;
import com.datastax.oss.driver.internal.core.specex.NoSpeculativeExecutionPolicy;
import com.datastax.oss.simulacron.common.cluster.ClusterSpec;
import com.datastax.oss.simulacron.common.stubbing.PrimeDsl;
import java.time.Duration;
import java.util.concurrent.TimeUnit;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.experimental.categories.Category;

@Category(ParallelizableTests.class)
public class SpeculativeExecutionIT {

  // Note: it looks like shorter delays cause precision issues with Netty timers
  private static final long SPECULATIVE_DELAY = 1000;

  private static final String QUERY_STRING = "select * from foo";
  private static final SimpleStatement QUERY = SimpleStatement.newInstance(QUERY_STRING);

  // Shared across all tests methods.
  @ClassRule
  public static final SimulacronRule SIMULACRON_RULE =
      new SimulacronRule(ClusterSpec.builder().withNodes(3));

  private final QueryCounter counter =
      QueryCounter.builder(SIMULACRON_RULE.cluster())
          .withFilter((l) -> l.getQuery().equals(QUERY_STRING))
          .build();

  @Before
  public void clear() {
    SIMULACRON_RULE.cluster().clearPrimes(true);
  }

  @Test
  public void should_not_start_speculative_executions_if_not_idempotent() {
    primeNode(
        0, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));

    try (CqlSession session = buildSession(2, SPECULATIVE_DELAY)) {
      ResultSet resultSet = session.execute(QUERY.setIdempotent(false));

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(0);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(0);

      counter.assertNodeCounts(1, 0, 0);
    }
  }

  @Test
  public void should_complete_from_first_speculative_execution_if_faster() {
    primeNode(
        0, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    primeNode(1, when(QUERY_STRING).then(noRows()));

    try (CqlSession session = buildSession(2, SPECULATIVE_DELAY)) {
      ResultSet resultSet = session.execute(QUERY);

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(1);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(1);

      counter.assertNodeCounts(1, 1, 0);
    }
  }

  @Test
  public void should_complete_from_initial_execution_if_speculative_is_started_but_slower() {
    primeNode(
        0, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    primeNode(
        1, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));

    try (CqlSession session = buildSession(2, SPECULATIVE_DELAY)) {
      ResultSet resultSet = session.execute(QUERY);

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(0);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(1);

      counter.assertNodeCounts(1, 1, 0);
    }
  }

  @Test
  public void should_complete_from_second_speculative_execution_if_faster() {
    primeNode(
        0, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    primeNode(
        1, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    primeNode(2, when(QUERY_STRING).then(noRows()));

    try (CqlSession session = buildSession(3, SPECULATIVE_DELAY)) {
      ResultSet resultSet = session.execute(QUERY);

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(2);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(2);

      counter.assertNodeCounts(1, 1, 1);
    }
  }

  @Test
  public void should_retry_within_initial_execution() {
    // This triggers a retry on the next node:
    primeNode(0, when(QUERY_STRING).then(isBootstrapping()));
    primeNode(1, when(QUERY_STRING).then(noRows()));

    try (CqlSession session = buildSession(3, SPECULATIVE_DELAY)) {
      ResultSet resultSet = session.execute(QUERY);

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(0);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(0);

      counter.assertNodeCounts(1, 1, 0);
    }
  }

  @Test
  public void should_retry_within_speculative_execution() {
    primeNode(
        0, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    // This triggers a retry on the next node:
    primeNode(1, when(QUERY_STRING).then(isBootstrapping()));
    primeNode(2, when(QUERY_STRING).then(noRows()));

    try (CqlSession session = buildSession(3, SPECULATIVE_DELAY)) {
      ResultSet resultSet = session.execute(QUERY);

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(1);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(1);

      counter.assertNodeCounts(1, 1, 1);
    }
  }

  @Test
  public void should_wait_for_last_execution_to_complete() {
    // Initial execution uses node0 which takes a long time to reply
    primeNode(
        0, when(QUERY_STRING).then(noRows()).delay(3 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    // specex1 runs fast, but only encounters failing nodes and stops
    primeNode(1, when(QUERY_STRING).then(isBootstrapping()));
    primeNode(2, when(QUERY_STRING).then(isBootstrapping()));

    try (CqlSession session = buildSession(2, SPECULATIVE_DELAY)) {
      ResultSet resultSet = session.execute(QUERY);

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(0);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(1);

      counter.assertNodeCounts(1, 1, 1);
    }
  }

  @Test(expected = AllNodesFailedException.class)
  public void should_fail_if_all_executions_reach_end_of_query_plan() {
    // Each execution gets a BOOTSTRAPPING response, but by the time it retries, the query plan will
    // be empty.
    for (int i = 0; i < 3; i++) {
      primeNode(
          i,
          when(QUERY_STRING)
              .then(isBootstrapping())
              .delay((3 - i) * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    }
    try (CqlSession session = buildSession(3, SPECULATIVE_DELAY)) {
      session.execute(QUERY);
    } finally {
      counter.assertNodeCounts(1, 1, 1);
    }
  }

  @Test
  public void should_allow_zero_delay() {
    // All executions start at the same time, but one of them is faster
    for (int i = 0; i < 2; i++) {
      primeNode(
          i, when(QUERY_STRING).then(noRows()).delay(2 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    }
    primeNode(2, when(QUERY_STRING).then(noRows()).delay(SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));

    try (CqlSession session = buildSession(3, 0)) {
      ResultSet resultSet = session.execute(QUERY);

      assertThat(resultSet.getExecutionInfo().getSuccessfulExecutionIndex()).isEqualTo(2);
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(2);

      counter.assertNodeCounts(1, 1, 1);
    }
  }

  @Test
  public void should_use_policy_from_request_profile() {
    // each host takes same amount of time
    for (int i = 0; i < 2; i++) {
      primeNode(
          i, when(QUERY_STRING).then(noRows()).delay(2 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    }

    // Set large delay for default so we ensure profile is used.
    try (CqlSession session = buildSessionWithProfile(3, 100, 2, 0)) {
      ResultSet resultSet = session.execute(QUERY.setExecutionProfileName("profile1"));

      // Expect only 1 speculative execution as that is all profile called for.
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(1);

      // Expect node 0 and 1 to be queried, but not 2.
      counter.assertNodeCounts(1, 1, 0);
    }
  }

  @Test
  public void should_use_policy_from_request_profile_when_not_configured_in_config() {
    // each host takes same amount of time
    for (int i = 0; i < 2; i++) {
      primeNode(
          i, when(QUERY_STRING).then(noRows()).delay(2 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    }

    // Disable in primary configuration
    try (CqlSession session = buildSessionWithProfile(-1, -1, 3, 0)) {
      ResultSet resultSet = session.execute(QUERY.setExecutionProfileName("profile1"));

      // Expect speculative executions on each node.
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(2);

      // Expect all nodes to be queried.
      counter.assertNodeCounts(1, 1, 1);
    }
  }

  @Test
  public void should_use_policy_from_config_when_not_configured_in_request_profile() {
    // each host takes same amount of time
    for (int i = 0; i < 2; i++) {
      primeNode(
          i, when(QUERY_STRING).then(noRows()).delay(2 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    }

    try (CqlSession session = buildSessionWithProfile(3, 0, 3, 0)) {
      // use profile where speculative execution is not configured.
      ResultSet resultSet = session.execute(QUERY.setExecutionProfileName("profile2"));

      // Expect speculative executions on each node since default configuration is used.
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(2);

      counter.assertNodeCounts(1, 1, 1);
    }
  }

  @Test
  public void should_not_speculatively_execute_when_defined_in_profile() {
    // each host takes same amount of time
    for (int i = 0; i < 2; i++) {
      primeNode(
          i, when(QUERY_STRING).then(noRows()).delay(2 * SPECULATIVE_DELAY, TimeUnit.MILLISECONDS));
    }

    // Disable in profile
    try (CqlSession session = buildSessionWithProfile(3, 100, -1, -1)) {
      ResultSet resultSet = session.execute(QUERY.setExecutionProfileName("profile1"));

      // Expect no speculative executions.
      assertThat(resultSet.getExecutionInfo().getSpeculativeExecutionCount()).isEqualTo(0);

      // Expect only node 0 to be queried since speculative execution is disabled for this profile.
      counter.assertNodeCounts(1, 0, 0);
    }
  }

  // Build a new Cluster instance for each test, because we need different configurations
  private CqlSession buildSession(int maxSpeculativeExecutions, long speculativeDelayMs) {
    return SessionUtils.newSession(
        SIMULACRON_RULE,
        SessionUtils.configLoaderBuilder()
            .withDuration(
                DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofMillis(SPECULATIVE_DELAY * 10))
            .withBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE, true)
            .withClass(
                DefaultDriverOption.LOAD_BALANCING_POLICY_CLASS, SortingLoadBalancingPolicy.class)
            .withClass(
                DefaultDriverOption.SPECULATIVE_EXECUTION_POLICY_CLASS,
                ConstantSpeculativeExecutionPolicy.class)
            .withInt(DefaultDriverOption.SPECULATIVE_EXECUTION_MAX, maxSpeculativeExecutions)
            .withDuration(
                DefaultDriverOption.SPECULATIVE_EXECUTION_DELAY,
                Duration.ofMillis(speculativeDelayMs))
            .build());
  }

  private CqlSession buildSessionWithProfile(
      int defaultMaxSpeculativeExecutions,
      long defaultSpeculativeDelayMs,
      int profile1MaxSpeculativeExecutions,
      long profile1SpeculativeDelayMs) {

    ProgrammaticDriverConfigLoaderBuilder builder =
        SessionUtils.configLoaderBuilder()
            .withDuration(
                DefaultDriverOption.REQUEST_TIMEOUT, Duration.ofMillis(SPECULATIVE_DELAY * 10))
            .withBoolean(DefaultDriverOption.REQUEST_DEFAULT_IDEMPOTENCE, true)
            .withClass(
                DefaultDriverOption.LOAD_BALANCING_POLICY_CLASS, SortingLoadBalancingPolicy.class);

    if (defaultMaxSpeculativeExecutions != -1 || defaultSpeculativeDelayMs != -1) {
      builder =
          builder.withClass(
              DefaultDriverOption.SPECULATIVE_EXECUTION_POLICY_CLASS,
              ConstantSpeculativeExecutionPolicy.class);
      if (defaultMaxSpeculativeExecutions != -1) {
        builder =
            builder.withInt(
                DefaultDriverOption.SPECULATIVE_EXECUTION_MAX, defaultMaxSpeculativeExecutions);
      }
      if (defaultSpeculativeDelayMs != -1) {
        builder =
            builder.withDuration(
                DefaultDriverOption.SPECULATIVE_EXECUTION_DELAY,
                Duration.ofMillis(defaultSpeculativeDelayMs));
      }
    } else {
      builder =
          builder.withClass(
              DefaultDriverOption.SPECULATIVE_EXECUTION_POLICY_CLASS,
              NoSpeculativeExecutionPolicy.class);
    }

    builder = builder.startProfile("profile1");
    if (profile1MaxSpeculativeExecutions != -1 || profile1SpeculativeDelayMs != -1) {
      builder =
          builder.withClass(
              DefaultDriverOption.SPECULATIVE_EXECUTION_POLICY_CLASS,
              ConstantSpeculativeExecutionPolicy.class);

      if (profile1MaxSpeculativeExecutions != -1) {
        builder =
            builder.withInt(
                DefaultDriverOption.SPECULATIVE_EXECUTION_MAX, profile1MaxSpeculativeExecutions);
      }
      if (profile1SpeculativeDelayMs != -1) {
        builder =
            builder.withDuration(
                DefaultDriverOption.SPECULATIVE_EXECUTION_DELAY,
                Duration.ofMillis(profile1SpeculativeDelayMs));
      }
    } else {
      builder =
          builder.withClass(
              DefaultDriverOption.SPECULATIVE_EXECUTION_POLICY_CLASS,
              NoSpeculativeExecutionPolicy.class);
    }

    builder =
        builder.startProfile("profile2").withString(DefaultDriverOption.REQUEST_CONSISTENCY, "ONE");

    CqlSession session = SessionUtils.newSession(SIMULACRON_RULE, builder.build());

    // validate profile data
    DriverContext context = session.getContext();
    DriverConfig driverConfig = context.getConfig();
    assertThat(driverConfig.getProfiles()).containsKeys("profile1", "profile2");

    assertThat(context.getSpeculativeExecutionPolicies())
        .hasSize(3)
        .containsKeys(DriverExecutionProfile.DEFAULT_NAME, "profile1", "profile2");

    SpeculativeExecutionPolicy defaultPolicy =
        context.getSpeculativeExecutionPolicy(DriverExecutionProfile.DEFAULT_NAME);
    SpeculativeExecutionPolicy policy1 = context.getSpeculativeExecutionPolicy("profile1");
    SpeculativeExecutionPolicy policy2 = context.getSpeculativeExecutionPolicy("profile2");
    Class<? extends SpeculativeExecutionPolicy> expectedDefaultPolicyClass =
        defaultMaxSpeculativeExecutions != -1 || defaultSpeculativeDelayMs != -1
            ? ConstantSpeculativeExecutionPolicy.class
            : NoSpeculativeExecutionPolicy.class;
    assertThat(defaultPolicy).isInstanceOf(expectedDefaultPolicyClass).isSameAs(policy2);

    // If configuration was same, same policy instance should be used.
    if (defaultMaxSpeculativeExecutions == profile1MaxSpeculativeExecutions
        && defaultSpeculativeDelayMs == profile1SpeculativeDelayMs) {
      assertThat(defaultPolicy).isSameAs(policy1);
    } else {
      assertThat(defaultPolicy).isNotSameAs(policy1);
    }

    Class<? extends SpeculativeExecutionPolicy> expectedProfile1PolicyClass =
        profile1MaxSpeculativeExecutions != -1 || profile1SpeculativeDelayMs != -1
            ? ConstantSpeculativeExecutionPolicy.class
            : NoSpeculativeExecutionPolicy.class;
    assertThat(policy1).isInstanceOf(expectedProfile1PolicyClass);

    return session;
  }

  private void primeNode(int id, PrimeDsl.PrimeBuilder primeBuilder) {
    SIMULACRON_RULE.cluster().node(id).prime(primeBuilder);
  }
}
