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

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.catchThrowable;
import static org.awaitility.Awaitility.await;

import com.datastax.oss.driver.api.core.ConsistencyLevel;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.DefaultConsistencyLevel;
import com.datastax.oss.driver.api.core.NoNodeAvailableException;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
import com.datastax.oss.driver.api.core.cql.SimpleStatement;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metadata.NodeState;
import com.datastax.oss.driver.api.core.metadata.TokenMap;
import com.datastax.oss.driver.api.testinfra.session.SessionUtils;
import com.datastax.oss.driver.api.testinfra.simulacron.SimulacronRule;
import com.datastax.oss.driver.categories.ParallelizableTests;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
import com.datastax.oss.protocol.internal.request.Query;
import com.datastax.oss.simulacron.common.cluster.ClusterSpec;
import com.datastax.oss.simulacron.common.cluster.QueryLog;
import com.datastax.oss.simulacron.common.stubbing.PrimeDsl;
import com.datastax.oss.simulacron.common.stubbing.PrimeDsl.RowBuilder;
import com.datastax.oss.simulacron.server.BoundCluster;
import com.datastax.oss.simulacron.server.BoundNode;
import com.datastax.oss.simulacron.server.BoundTopic;
import com.tngtech.java.junit.dataprovider.DataProvider;
import com.tngtech.java.junit.dataprovider.DataProviderRunner;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.runner.RunWith;

@Category(ParallelizableTests.class)
@RunWith(DataProviderRunner.class)
public class AllLoadBalancingPoliciesSimulacronIT {

  @ClassRule
  public static final SimulacronRule SIMULACRON_RULE =
      new SimulacronRule(ClusterSpec.builder().withNodes(5, 5, 5));

  @Before
  public void reset() {
    SIMULACRON_RULE.cluster().start();
    SIMULACRON_RULE.cluster().clearLogs();
    SIMULACRON_RULE.cluster().clearPrimes(true);
    SIMULACRON_RULE
        .cluster()
        .prime(
            PrimeDsl.when("SELECT * FROM system_schema.keyspaces")
                .then(new RowBuilder().columnTypes(KEYSPACE_COLUMNS).row(KEYSPACE_ROW).build()));
  }

  @Test
  @DataProvider({
    "BasicLoadBalancingPolicy,dc1",
    "DefaultLoadBalancingPolicy,dc1",
    "DcInferringLoadBalancingPolicy,dc1",
    "DcInferringLoadBalancingPolicy,null",
  })
  public void should_round_robin_within_local_dc_when_dc_aware_but_not_token_aware(
      String lbp, String dc) {

    // given: DC is provided or inferred, token awareness is disabled and remote DCs are allowed
    try (CqlSession session = newSession(lbp, dc, 2, true, false)) {

      // when: a query is executed 50 times.
      for (int i = 0; i < 50; i++) {
        session.execute(QUERY);
      }

      // then: each node in local DC should get an equal number of requests.
      for (int i = 0; i < 5; i++) {
        assertThat(queries(0, i).count()).isEqualTo(10);
      }

      // then: no node in the remote DC should get a request.
      assertThat(queries(1).count()).isEqualTo(0);
      assertThat(queries(2).count()).isEqualTo(0);
    }
  }

  @Test
  @DataProvider({
    "BasicLoadBalancingPolicy,dc1,ONE",
    "BasicLoadBalancingPolicy,dc1,LOCAL_ONE",
    "BasicLoadBalancingPolicy,dc1,TWO",
    "BasicLoadBalancingPolicy,dc1,QUORUM",
    "BasicLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DefaultLoadBalancingPolicy,dc1,ONE",
    "DefaultLoadBalancingPolicy,dc1,LOCAL_ONE",
    "DefaultLoadBalancingPolicy,dc1,TWO",
    "DefaultLoadBalancingPolicy,dc1,QUORUM",
    "DefaultLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DcInferringLoadBalancingPolicy,dc1,ONE",
    "DcInferringLoadBalancingPolicy,dc1,LOCAL_ONE",
    "DcInferringLoadBalancingPolicy,dc1,TWO",
    "DcInferringLoadBalancingPolicy,dc1,QUORUM",
    "DcInferringLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DcInferringLoadBalancingPolicy,null,ONE",
    "DcInferringLoadBalancingPolicy,null,LOCAL_ONE",
    "DcInferringLoadBalancingPolicy,null,TWO",
    "DcInferringLoadBalancingPolicy,null,QUORUM",
    "DcInferringLoadBalancingPolicy,null,LOCAL_QUORUM",
  })
  public void should_use_local_replicas_when_dc_aware_and_token_aware_and_enough_local_replicas_up(
      String lbp, String dc, DefaultConsistencyLevel cl) {

    // given: DC is provided or inferred, token awareness enabled, remotes allowed, CL <= 2
    try (CqlSession session = newSession(lbp, dc, 2, true)) {

      // given: one replica and 2 non-replicas down in local DC, but CL <= 2 still achievable
      List<Node> aliveReplicas = degradeLocalDc(session);

      // when: a query is executed 50 times and some nodes are down in the local DC.
      for (int i = 0; i < 50; i++) {
        session.execute(
            SimpleStatement.newInstance(QUERY)
                .setConsistencyLevel(cl)
                .setRoutingKeyspace("test")
                .setRoutingKey(ROUTING_KEY));
      }

      // then: all requests should be distributed to the remaining up replicas in local DC
      BoundNode alive1 = findNode(aliveReplicas.get(0));
      BoundNode alive2 = findNode(aliveReplicas.get(1));
      assertThat(queries(alive1).count() + queries(alive2).count()).isEqualTo(50);

      // then: no node in the remote DCs should get a request.
      assertThat(queries(1).count()).isEqualTo(0);
      assertThat(queries(2).count()).isEqualTo(0);
    }
  }

  @Test
  public void should_round_robin_within_all_dcs_when_dc_agnostic() {

    // given: DC-agnostic LBP, no local DC, remotes not allowed, token awareness enabled
    try (CqlSession session = newSession("BasicLoadBalancingPolicy", null, 0, false)) {

      // when: a query is executed 150 times.
      for (int i = 0; i < 150; i++) {
        session.execute(
            SimpleStatement.newInstance(QUERY)
                // local CL should be ignored since there is no local DC
                .setConsistencyLevel(ConsistencyLevel.LOCAL_QUORUM));
      }

      // then: each node should get 10 requests, even remote ones since the LBP is DC-agnostic.
      for (int dc = 0; dc < 3; dc++) {
        for (int n = 0; n < 5; n++) {
          assertThat(queries(dc, n).count()).isEqualTo(10);
        }
      }
    }
  }

  @Test
  @DataProvider({
    "BasicLoadBalancingPolicy,dc1,ONE",
    "BasicLoadBalancingPolicy,dc1,TWO",
    "BasicLoadBalancingPolicy,dc1,THREE",
    "BasicLoadBalancingPolicy,dc1,QUORUM",
    "BasicLoadBalancingPolicy,dc1,ANY",
    "DefaultLoadBalancingPolicy,dc1,ONE",
    "DefaultLoadBalancingPolicy,dc1,TWO",
    "DefaultLoadBalancingPolicy,dc1,THREE",
    "DefaultLoadBalancingPolicy,dc1,QUORUM",
    "DefaultLoadBalancingPolicy,dc1,ANY",
    "DcInferringLoadBalancingPolicy,dc1,ONE",
    "DcInferringLoadBalancingPolicy,dc1,TWO",
    "DcInferringLoadBalancingPolicy,dc1,THREE",
    "DcInferringLoadBalancingPolicy,dc1,QUORUM",
    "DcInferringLoadBalancingPolicy,dc1,ANY",
    "DcInferringLoadBalancingPolicy,null,ONE",
    "DcInferringLoadBalancingPolicy,null,TWO",
    "DcInferringLoadBalancingPolicy,null,THREE",
    "DcInferringLoadBalancingPolicy,null,QUORUM",
    "DcInferringLoadBalancingPolicy,null,ANY",
  })
  public void should_use_remote_nodes_when_no_up_nodes_in_local_dc_for_non_local_cl(
      String lbp, String dc, DefaultConsistencyLevel cl) {

    // given: 1 remote allowed per DC and a non-local CL, token awareness enabled
    try (CqlSession session = newSession(lbp, dc, 1, false)) {

      // given: local DC is down
      stopLocalDc(session);

      // when: a query is executed 50 times and all nodes are down in local DC.
      for (int i = 0; i < 50; i++) {
        session.execute(
            SimpleStatement.newInstance(QUERY)
                .setConsistencyLevel(cl)
                .setRoutingKeyspace("test")
                .setRoutingKey(ROUTING_KEY));
      }

      // then: only 1 node in each remote DC should get requests (we can't know which ones exactly).
      assertThat(queries(1).count() + queries(2).count()).isEqualTo(50);
    }
  }

  @Test
  @DataProvider({
    "BasicLoadBalancingPolicy,dc1,LOCAL_ONE",
    "BasicLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "BasicLoadBalancingPolicy,dc1,LOCAL_SERIAL",
    "DefaultLoadBalancingPolicy,dc1,LOCAL_ONE",
    "DefaultLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DefaultLoadBalancingPolicy,dc1,LOCAL_SERIAL",
    "DcInferringLoadBalancingPolicy,dc1,LOCAL_ONE",
    "DcInferringLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DcInferringLoadBalancingPolicy,dc1,LOCAL_SERIAL",
    "DcInferringLoadBalancingPolicy,null,LOCAL_ONE",
    "DcInferringLoadBalancingPolicy,null,LOCAL_QUORUM",
    "DcInferringLoadBalancingPolicy,null,LOCAL_SERIAL",
  })
  public void should_not_use_remote_nodes_when_using_local_cl(
      String lbp, String dc, DefaultConsistencyLevel cl) {

    // given: remotes allowed but not for local CL, token awareness enabled, local CL
    try (CqlSession session = newSession(lbp, dc, 5, false)) {

      // given: local DC is down
      stopLocalDc(session);

      // when: a query is executed 50 times and all nodes are down in local DC.
      for (int i = 0; i < 50; i++) {
        Throwable t =
            catchThrowable(
                () ->
                    session.execute(
                        SimpleStatement.newInstance(QUERY)
                            .setConsistencyLevel(cl)
                            .setRoutingKeyspace("test")
                            .setRoutingKey(ROUTING_KEY)));

        // then: expect a NNAE for a local CL since no local replicas available.
        assertThat(t).isInstanceOf(NoNodeAvailableException.class);
      }

      // then: no node in the remote DCs should get a request.
      assertThat(queries(1).count()).isEqualTo(0);
      assertThat(queries(2).count()).isEqualTo(0);
    }
  }

  @Test
  @DataProvider({
    "BasicLoadBalancingPolicy,dc1,LOCAL_ONE",
    "BasicLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DefaultLoadBalancingPolicy,dc1,LOCAL_ONE",
    "DefaultLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DcInferringLoadBalancingPolicy,dc1,LOCAL_ONE",
    "DcInferringLoadBalancingPolicy,dc1,LOCAL_QUORUM",
    "DcInferringLoadBalancingPolicy,null,LOCAL_ONE",
    "DcInferringLoadBalancingPolicy,null,LOCAL_QUORUM",
  })
  public void should_use_remote_nodes_when_using_local_cl_if_allowed(
      String lbp, String dc, DefaultConsistencyLevel cl) {

    // given: only one node allowed per remote DC and remotes allowed even for local CLs.
    try (CqlSession session = newSession(lbp, dc, 1, true)) {

      // given: local DC is down
      stopLocalDc(session);

      // when: a query is executed 50 times and all nodes are down in local DC.
      for (int i = 0; i < 50; i++) {
        session.execute(
            SimpleStatement.newInstance(QUERY)
                .setConsistencyLevel(cl)
                .setRoutingKeyspace("test")
                .setRoutingKey(ROUTING_KEY));
      }

      // then: only 1 node in each remote DC should get requests (we can't know which ones exactly).
      assertThat(queries(1).count() + queries(2).count()).isEqualTo(50);
    }
  }

  @Test
  @DataProvider({
    "BasicLoadBalancingPolicy,dc1",
    "DefaultLoadBalancingPolicy,dc1",
    "DcInferringLoadBalancingPolicy,dc1",
    "DcInferringLoadBalancingPolicy,null"
  })
  public void should_not_use_excluded_dc_using_node_filter(String lbp, String dc) {

    // given: remotes allowed even for local CLs, but node filter excluding dc2
    try (CqlSession session = newSession(lbp, dc, 5, true, true, excludeDc("dc2"))) {

      // when: A query is made and nodes for the local dc are available.
      for (int i = 0; i < 50; i++) {
        session.execute(
            SimpleStatement.newInstance(QUERY)
                .setRoutingKeyspace("test")
                .setRoutingKey(ROUTING_KEY));
      }

      // then: only nodes in the local DC should have been queried.
      assertThat(queries(0).count()).isEqualTo(50);
      assertThat(queries(1).count()).isEqualTo(0);
      assertThat(queries(2).count()).isEqualTo(0);

      // given: local DC is down
      stopLocalDc(session);

      SIMULACRON_RULE.cluster().clearLogs();

      // when: A query is made and all nodes in the local dc are down.
      for (int i = 0; i < 50; i++) {
        session.execute(
            SimpleStatement.newInstance(QUERY)
                .setRoutingKeyspace("test")
                .setRoutingKey(ROUTING_KEY));
      }

      // then: Only nodes in DC3 should have been queried, since DC2 is excluded and DC1 is down.
      assertThat(queries(0).count()).isEqualTo(0);
      assertThat(queries(1).count()).isEqualTo(0);
      assertThat(queries(2).count()).isEqualTo(50);
    }
  }

  private static final ByteBuffer ROUTING_KEY = ByteBuffer.wrap(new byte[] {1, 2, 3, 4});

  private static final String[] KEYSPACE_COLUMNS =
      new String[] {
        "keyspace_name", "varchar",
        "durable_writes", "boolean",
        "replication", "map<varchar, varchar>"
      };

  private static final Object[] KEYSPACE_ROW =
      new Object[] {
        "keyspace_name",
        "test",
        "durable_writes",
        true,
        "replication",
        ImmutableMap.of(
            "class",
            "org.apache.cassandra.locator.NetworkTopologyStrategy",
            "dc1",
            "3",
            "dc2",
            "3",
            "dc3",
            "3")
      };

  private static final String QUERY = "SELECT * FROM test.foo";

  private CqlSession newSession(String lbp, String dc, int maxRemoteNodes, boolean allowLocalCl) {
    return newSession(lbp, dc, maxRemoteNodes, allowLocalCl, true);
  }

  private CqlSession newSession(
      String lbp, String dc, int maxRemoteNodes, boolean allowLocalCl, boolean tokenAware) {
    return newSession(lbp, dc, maxRemoteNodes, allowLocalCl, tokenAware, null);
  }

  private CqlSession newSession(
      String lbp,
      String dc,
      int maxRemoteNodes,
      boolean allowLocalCl,
      boolean tokenAware,
      Predicate<Node> nodeFilter) {
    DriverConfigLoader loader =
        SessionUtils.configLoaderBuilder()
            .withBoolean(DefaultDriverOption.METADATA_SCHEMA_ENABLED, tokenAware)
            .withString(DefaultDriverOption.LOAD_BALANCING_POLICY_CLASS, lbp)
            .withString(DefaultDriverOption.LOAD_BALANCING_LOCAL_DATACENTER, dc)
            .withInt(
                DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_MAX_NODES_PER_REMOTE_DC,
                maxRemoteNodes)
            .withBoolean(
                DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_ALLOW_FOR_LOCAL_CONSISTENCY_LEVELS,
                allowLocalCl)
            .build();
    return SessionUtils.newSession(SIMULACRON_RULE, null, null, null, nodeFilter, loader);
  }

  private BoundNode findNode(Node node) {
    BoundCluster simulacron = SIMULACRON_RULE.cluster();
    SocketAddress toFind = node.getEndPoint().resolve();
    for (BoundNode boundNode : simulacron.getNodes()) {
      if (boundNode.getAddress().equals(toFind)) {
        return boundNode;
      }
    }
    throw new AssertionError("Could not find node: " + toFind);
  }

  private void stopLocalDc(CqlSession session) {
    SIMULACRON_RULE.cluster().dc(0).stop();
    awaitDown(nodesInDc(session, "dc1"));
  }

  private List<Node> degradeLocalDc(CqlSession session) {
    // stop 1 replica and 2 non-replicas in dc1
    List<Node> localReplicas = replicasInDc(session, "dc1");
    assertThat(localReplicas).hasSize(3);
    BoundNode replica1 = findNode(localReplicas.get(0));

    List<Node> localOthers = nonReplicasInDc(session, "dc1");
    assertThat(localOthers).hasSize(2);
    BoundNode other1 = findNode(localOthers.get(0));
    BoundNode other2 = findNode(localOthers.get(1));

    replica1.stop();
    other1.stop();
    other2.stop();

    awaitDown(localReplicas.get(0), localOthers.get(0), localOthers.get(1));
    return localReplicas.subList(1, 3);
  }

  private Stream<QueryLog> queries(int dc, int node) {
    return queries(SIMULACRON_RULE.cluster().dc(dc).node(node));
  }

  private Stream<QueryLog> queries(int dc) {
    return queries(SIMULACRON_RULE.cluster().dc(dc));
  }

  private Stream<QueryLog> queries(BoundTopic<?, ?> topic) {
    return topic.getLogs().getQueryLogs().stream()
        .filter(q -> q.getFrame().message instanceof Query)
        .filter(q -> ((Query) q.getFrame().message).query.equals(QUERY));
  }

  private List<Node> nodesInDc(CqlSession session, String dcName) {
    return session.getMetadata().getNodes().values().stream()
        .filter(n -> Objects.equals(n.getDatacenter(), dcName))
        .collect(Collectors.toList());
  }

  private List<Node> replicasInDc(CqlSession session, String dcName) {
    assertThat(session.getMetadata().getTokenMap()).isPresent();
    TokenMap tokenMap = session.getMetadata().getTokenMap().get();
    return tokenMap.getReplicas("test", ROUTING_KEY).stream()
        .filter(n -> Objects.equals(n.getDatacenter(), dcName))
        .collect(Collectors.toList());
  }

  private List<Node> nonReplicasInDc(
      CqlSession session, @SuppressWarnings("SameParameterValue") String dcName) {
    List<Node> nodes = nodesInDc(session, dcName);
    nodes.removeAll(replicasInDc(session, dcName));
    return nodes;
  }

  private Predicate<Node> excludeDc(@SuppressWarnings("SameParameterValue") String dcName) {
    return node -> !Objects.equals(node.getDatacenter(), dcName);
  }

  private void awaitDown(Node... nodes) {
    awaitDown(Arrays.asList(nodes));
  }

  private void awaitDown(Iterable<Node> nodes) {
    await()
        .atMost(Duration.ofSeconds(10))
        .untilAsserted(
            () -> {
              for (Node node : nodes) {
                assertThat(node.getState()).isEqualTo(NodeState.DOWN);
              }
            });
  }
}
