/*
 * 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 com.datastax.driver.core.Assertions.assertThat;

import com.datastax.driver.core.exceptions.InvalidTypeException;
import com.datastax.driver.core.policies.LoadBalancingPolicy;
import com.datastax.driver.core.policies.RoundRobinPolicy;
import com.datastax.driver.core.policies.WhiteListPolicy;
import com.datastax.driver.core.utils.CassandraVersion;
import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import org.testng.annotations.Test;

/**
 * This class uses subclasses for each type of partitioner.
 *
 * <p>There's normally a way to parametrize a TestNG class with @Factory and @DataProvider, but it
 * doesn't seem to work with multiple methods.
 */
@CCMConfig(numberOfNodes = 3, createKeyspace = false)
public abstract class TokenIntegrationTest extends CCMTestsSupport {

  private final DataType expectedTokenType;
  private final int numTokens;
  private final boolean useVnodes;
  private String ks1;
  private String ks2;

  public TokenIntegrationTest(DataType expectedTokenType, boolean useVnodes) {
    this.expectedTokenType = expectedTokenType;
    this.numTokens = useVnodes ? 256 : 1;
    this.useVnodes = useVnodes;
  }

  @Override
  public Cluster.Builder createClusterBuilder() {
    // Only connect to node 1, which makes it easier to query system tables in
    // should_expose_tokens_per_host()
    LoadBalancingPolicy lbp =
        new WhiteListPolicy(new RoundRobinPolicy(), Collections.singleton(ccm().addressOfNode(1)));
    return Cluster.builder()
        .addContactPoints(getContactPoints().get(0))
        .withPort(ccm().getBinaryPort())
        .withLoadBalancingPolicy(lbp);
  }

  @Override
  public void onTestContextInitialized() {
    ks1 = TestUtils.generateIdentifier("ks_");
    ks2 = TestUtils.generateIdentifier("ks_");
    execute(
        String.format(
            "CREATE KEYSPACE %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 1}",
            ks1),
        String.format(
            "CREATE KEYSPACE %s WITH replication = {'class': 'SimpleStrategy', 'replication_factor': 2}",
            ks2),
        String.format("USE %s", ks1),
        "CREATE TABLE foo(i int primary key)",
        "INSERT INTO foo (i) VALUES (1)",
        "INSERT INTO foo (i) VALUES (2)",
        "INSERT INTO foo (i) VALUES (3)");
  }

  /**
   * Validates that {@link TokenRange}s are exposed via a {@link Cluster}'s {@link Metadata} and
   * they can be used to query data.
   *
   * @test_category metadata:token
   * @expected_result token ranges are exposed and usable.
   * @jira_ticket JAVA-312
   * @since 2.0.10, 2.1.5
   */
  @Test(groups = "short")
  public void should_expose_token_ranges() throws Exception {
    Metadata metadata = cluster().getMetadata();

    // Find the replica for a given partition key
    int testKey = 1;
    Set<Host> replicas =
        metadata.getReplicas(
            ks1,
            TypeCodec.cint()
                .serialize(
                    testKey,
                    cluster().getConfiguration().getProtocolOptions().getProtocolVersion()));
    assertThat(replicas).hasSize(1);
    Host replica = replicas.iterator().next();

    // Iterate the cluster's token ranges. For each one, use a range query to ask Cassandra which
    // partition keys
    // are in this range.

    PreparedStatement rangeStmt =
        session().prepare("SELECT i FROM foo WHERE token(i) > ? and token(i) <= ?");

    TokenRange foundRange = null;
    for (TokenRange range : metadata.getTokenRanges()) {
      List<Row> rows = rangeQuery(rangeStmt, range);
      for (Row row : rows) {
        if (row.getInt("i") == testKey) {
          // We should find our test key exactly once
          assertThat(foundRange)
              .describedAs("found the same key in two ranges: " + foundRange + " and " + range)
              .isNull();
          foundRange = range;
          // That range should be managed by the replica
          assertThat(metadata.getReplicas(ks1, range)).contains(replica);
        }
      }
    }
    assertThat(foundRange).isNotNull();
  }

  private List<Row> rangeQuery(PreparedStatement rangeStmt, TokenRange range) {
    List<Row> rows = Lists.newArrayList();
    for (TokenRange subRange : range.unwrap()) {
      Statement statement = rangeStmt.bind(subRange.getStart(), subRange.getEnd());
      rows.addAll(session().execute(statement).all());
    }
    return rows;
  }

  /**
   * Validates that a {@link Token} can be retrieved and parsed by executing 'select token(name)'
   * and then used to find data matching that token.
   *
   * <p>
   *
   * <p>This test does the following:
   *
   * <p>
   *
   * <ol>
   *   <li>Retrieve the token for the key with value '1', get it by index, and ensure if is of the
   *       expected token type.
   *   <li>Retrieve the token for the partition key with getPartitionKeyToken
   *   <li>Select data by token with a BoundStatement.
   *   <li>Select data by token using setToken by index.
   *   <li>Select data by token with setPartitionKeyToken.
   * </ol>
   *
   * @test_category token
   * @expected_result tokens are selectable, properly parsed, and usable as input.
   * @jira_ticket JAVA-312
   * @since 2.0.10, 2.1.5
   */
  @Test(groups = "short")
  public void should_get_token_from_row_and_set_token_in_query() {
    // get by index:
    Row row = session().execute("SELECT token(i) FROM foo WHERE i = 1").one();
    Token token = row.getToken(0);
    assertThat(token.getType()).isEqualTo(expectedTokenType);

    assertThat(row.getPartitionKeyToken()).isEqualTo(token);

    PreparedStatement pst = session().prepare("SELECT * FROM foo WHERE token(i) = ?");
    row = session().execute(pst.bind(token)).one();
    assertThat(row.getInt(0)).isEqualTo(1);

    row = session().execute(pst.bind().setToken(0, token)).one();
    assertThat(row.getInt(0)).isEqualTo(1);

    row = session().execute(pst.bind().setPartitionKeyToken(token)).one();
    assertThat(row.getInt(0)).isEqualTo(1);
  }

  /**
   * Validates that a {@link Token} can be retrieved and parsed by using bind variables and
   * aliasing.
   *
   * <p>
   *
   * <p>This test does the following:
   *
   * <p>
   *
   * <ol>
   *   <li>Retrieve the token by alias for the key '1', and ensure it matches the token by index.
   *   <li>Select data by token using setToken by name.
   * </ol>
   */
  @Test(groups = "short")
  @CassandraVersion("2.0")
  public void should_get_token_from_row_and_set_token_in_query_with_binding_and_aliasing() {
    Row row = session().execute("SELECT token(i) AS t FROM foo WHERE i = 1").one();
    Token token = row.getToken("t");
    assertThat(token.getType()).isEqualTo(expectedTokenType);

    PreparedStatement pst = session().prepare("SELECT * FROM foo WHERE token(i) = :myToken");
    row = session().execute(pst.bind().setToken("myToken", token)).one();
    assertThat(row.getInt(0)).isEqualTo(1);

    row = session().execute("SELECT * FROM foo WHERE token(i) = ?", token).one();
    assertThat(row.getInt(0)).isEqualTo(1);
  }

  /**
   * Ensures that an exception is raised when attempting to retrieve a token a non-token column.
   *
   * @test_category token
   * @expected_result an exception is raised.
   * @jira_ticket JAVA-312
   * @since 2.0.10, 2.1.5
   */
  @Test(groups = "short", expectedExceptions = InvalidTypeException.class)
  public void should_raise_exception_when_get_token_on_non_token() {
    Row row = session().execute("SELECT i FROM foo WHERE i = 1").one();
    row.getToken(0);
  }

  /**
   * Ensures that @{link TokenRange}s are exposed at a per host level, the ranges are complete, the
   * entire ring is represented, and that ranges do not overlap.
   *
   * <p>
   *
   * <p>Also ensures that ranges from another replica are present when a Host is a replica for
   * another node.
   *
   * @test_category metadata:token
   * @expected_result The entire token range is represented collectively and the ranges do not
   *     overlap.
   * @jira_ticket JAVA-312
   * @since 2.0.10, 2.1.5
   */
  @Test(groups = "short")
  public void should_expose_token_ranges_per_host() {
    checkRangesPerHost(ks1, 1);
    checkRangesPerHost(ks2, 2);
    assertThat(cluster()).hasValidTokenRanges();
  }

  private void checkRangesPerHost(String keyspace, int replicationFactor) {
    List<TokenRange> allRangesWithReplicas = Lists.newArrayList();

    // Get each host's ranges, the count should match the replication factor
    for (int i = 1; i <= 3; i++) {
      Host host = TestUtils.findHost(cluster(), i);
      Set<TokenRange> hostRanges = cluster().getMetadata().getTokenRanges(keyspace, host);
      // Special case: When using vnodes the tokens are not evenly assigned to each replica.
      if (!useVnodes) {
        assertThat(hostRanges).hasSize(replicationFactor * numTokens);
      }
      allRangesWithReplicas.addAll(hostRanges);
    }

    // Special case check for vnodes to ensure that total number of replicated ranges is correct.
    assertThat(allRangesWithReplicas).hasSize(3 * numTokens * replicationFactor);

    // Once we ignore duplicates, the number of ranges should match the number of nodes.
    Set<TokenRange> allRanges = new HashSet<TokenRange>(allRangesWithReplicas);
    assertThat(allRanges).hasSize(3 * numTokens);

    // And the ranges should cover the whole ring and no ranges intersect.
    assertThat(cluster()).hasValidTokenRanges(keyspace);
  }

  /**
   * Ensures that Tokens are exposed for each Host and that the match those in the system tables.
   *
   * <p>
   *
   * <p>Also validates that tokens are not present for multiple hosts.
   *
   * @test_category metadata:token
   * @expected_result Tokens are exposed by Host and match those in the system tables.
   * @jira_ticket JAVA-312
   * @since 2.0.10, 2.1.5
   */
  @Test(groups = "short")
  public void should_expose_tokens_per_host() {
    for (Host host : cluster().getMetadata().allHosts()) {
      assertThat(host.getTokens()).hasSize(numTokens);
      // Check against the info in the system tables, which is a bit weak since it's exactly how the
      // metadata is
      // constructed in the first place, but there's not much else we can do.
      // Note that this relies on all queries going to node 1, which is why we use a WhiteList LBP
      // in setup().
      boolean isControlHost =
          host.getEndPoint()
              .equals(cluster().manager.controlConnection.connectionRef.get().endPoint);
      Row row;
      if (isControlHost) {
        row = session().execute("select tokens from system.local").one();
      } else {
        // non-control hosts are populated from system.peers and their broadcast address should be
        // known
        assertThat(host.getBroadcastSocketAddress()).isNotNull();
        row =
            session()
                .execute(
                    "select tokens from system.peers where peer = '"
                        + host.getBroadcastSocketAddress().getAddress().getHostAddress()
                        + "'")
                .one();
      }
      Set<String> tokenStrings = row.getSet("tokens", String.class);
      assertThat(tokenStrings).hasSize(numTokens);
      Iterable<Token> tokensFromSystemTable =
          Iterables.transform(
              tokenStrings,
              new Function<String, Token>() {
                @Override
                public Token apply(String input) {
                  return tokenFactory().fromString(input);
                }
              });

      assertThat(host.getTokens())
          .containsOnlyOnce(Iterables.toArray(tokensFromSystemTable, Token.class));
    }
  }

  /**
   * Ensures that for the {@link TokenRange}s returned by {@link Metadata#getTokenRanges()} that
   * there exists at most one {@link TokenRange} for which calling {@link
   * TokenRange#isWrappedAround()} returns true and {@link TokenRange#unwrap()} returns two {@link
   * TokenRange}s.
   *
   * @test_category metadata:token
   * @expected_result Tokens are exposed by Host and match those in the system tables.
   * @jira_ticket JAVA-312
   * @since 2.0.10, 2.1.5
   */
  @Test(groups = "short")
  public void should_only_unwrap_one_range_for_all_ranges() {
    Set<TokenRange> ranges = cluster().getMetadata().getTokenRanges();

    assertOnlyOneWrapped(ranges);

    Iterable<TokenRange> splitRanges =
        Iterables.concat(
            Iterables.transform(
                ranges,
                new Function<TokenRange, Iterable<TokenRange>>() {
                  @Override
                  public Iterable<TokenRange> apply(TokenRange input) {
                    return input.splitEvenly(10);
                  }
                }));

    assertOnlyOneWrapped(splitRanges);
  }

  /**
   * Asserts that given the input {@link TokenRange}s that at most one of them wraps the token ring.
   *
   * @param ranges Ranges to validate against.
   */
  protected void assertOnlyOneWrapped(Iterable<TokenRange> ranges) {
    TokenRange wrappedRange = null;

    for (TokenRange range : ranges) {
      if (range.isWrappedAround()) {
        assertThat(wrappedRange)
            .as(
                "Found a wrapped around TokenRange (%s) when one already exists (%s).",
                range, wrappedRange)
            .isNull();
        wrappedRange = range;

        assertThat(range).isWrappedAround(); // this also checks the unwrapped ranges
      } else {
        assertThat(range).isNotWrappedAround();
      }
    }
  }

  @Test(groups = "short")
  public void should_expose_token_and_range_creation_methods() {
    Metadata metadata = cluster().getMetadata();

    // Pick a random range
    TokenRange range = metadata.getTokenRanges().iterator().next();

    Token start = metadata.newToken(range.getStart().toString());
    Token end = metadata.newToken(range.getEnd().toString());

    assertThat(metadata.newTokenRange(start, end)).isEqualTo(range);
  }

  @Test(groups = "short")
  public void should_create_token_from_partition_key() {
    Metadata metadata = cluster().getMetadata();

    Row row = session().execute("SELECT token(i) FROM foo WHERE i = 1").one();
    Token expected = row.getToken(0);

    ProtocolVersion protocolVersion =
        cluster().getConfiguration().getProtocolOptions().getProtocolVersion();
    assertThat(metadata.newToken(TypeCodec.cint().serialize(1, protocolVersion)))
        .isEqualTo(expected);
  }

  protected abstract Token.Factory tokenFactory();
}
