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

import static com.datastax.driver.core.querybuilder.QueryBuilder.batch;
import static com.datastax.driver.core.querybuilder.QueryBuilder.eq;
import static com.datastax.driver.core.querybuilder.QueryBuilder.insertInto;
import static com.datastax.driver.core.querybuilder.QueryBuilder.select;
import static com.datastax.driver.core.querybuilder.QueryBuilder.timestamp;
import static com.datastax.driver.core.querybuilder.QueryBuilder.ttl;
import static com.datastax.driver.core.querybuilder.QueryBuilder.update;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;

import com.datastax.driver.core.CCMConfig;
import com.datastax.driver.core.CCMTestsSupport;
import com.datastax.driver.core.CodecRegistry;
import com.datastax.driver.core.ProtocolVersion;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.TableMetadata;
import java.nio.ByteBuffer;
import org.testng.annotations.Test;

@CCMConfig(clusterProvider = "createClusterBuilderNoDebouncing")
public class QueryBuilderRoutingKeyTest extends CCMTestsSupport {

  private static final String TABLE_TEXT = "test_text";
  private static final String TABLE_INT = "test_int";
  private static final String TABLE_CASE = "test_case";
  private static final String TABLE_CASE_QUOTED = "test_case_quoted";

  @Override
  public void onTestContextInitialized() {
    execute(
        String.format("CREATE TABLE %s (k text PRIMARY KEY, a int, b int)", TABLE_TEXT),
        String.format("CREATE TABLE %s (k int PRIMARY KEY, a int, b int)", TABLE_INT),
        String.format("CREATE TABLE %s (theKey int PRIMARY KEY, a int, b int)", TABLE_CASE),
        String.format(
            "CREATE TABLE %s (\"theKey\" int PRIMARY KEY, a int, b int, \"tHEkEY\" int)",
            TABLE_CASE_QUOTED));
  }

  @Test(groups = "short")
  public void textRoutingKeyTest() throws Exception {

    BuiltStatement query;
    TableMetadata table = cluster().getMetadata().getKeyspace(keyspace).getTable(TABLE_TEXT);
    assertNotNull(table);
    ProtocolVersion protocolVersion =
        cluster().getConfiguration().getProtocolOptions().getProtocolVersion();
    CodecRegistry codecRegistry = CodecRegistry.DEFAULT_INSTANCE;

    String txt = "If she weighs the same as a duck... she's made of wood.";
    query = insertInto(table).values(new String[] {"k", "a", "b"}, new Object[] {txt, 1, 2});
    assertEquals(
        query.getRoutingKey(protocolVersion, codecRegistry), ByteBuffer.wrap(txt.getBytes()));
    session().execute(query);

    query = select().from(table).where(eq("k", txt));
    assertEquals(
        query.getRoutingKey(protocolVersion, codecRegistry), ByteBuffer.wrap(txt.getBytes()));
    Row row = session().execute(query).one();
    assertEquals(row.getString("k"), txt);
    assertEquals(row.getInt("a"), 1);
    assertEquals(row.getInt("b"), 2);
  }

  @Test(groups = "short")
  public void routingKeyColumnCaseSensitivityTest() throws Exception {

    BuiltStatement query;
    TableMetadata table = cluster().getMetadata().getKeyspace(keyspace).getTable(TABLE_CASE);
    assertNotNull(table);
    ProtocolVersion protocolVersion =
        cluster().getConfiguration().getProtocolOptions().getProtocolVersion();
    CodecRegistry codecRegistry = CodecRegistry.DEFAULT_INSTANCE;

    query = insertInto(table).values(new String[] {"theKey", "a", "b"}, new Object[] {42, 1, 2});
    ByteBuffer bb = ByteBuffer.allocate(4);
    bb.putInt(0, 42);
    assertEquals(query.getRoutingKey(protocolVersion, codecRegistry), bb);
    session().execute(query);

    query = select().from(table).where(eq("theKey", 42));
    assertEquals(query.getRoutingKey(protocolVersion, codecRegistry), bb);
    Row row = session().execute(query).one();
    assertEquals(row.getInt("theKey"), 42);
    assertEquals(row.getInt("a"), 1);
    assertEquals(row.getInt("b"), 2);

    query = insertInto(table).values(new String[] {"ThEkEy", "a", "b"}, new Object[] {42, 1, 2});
    bb = ByteBuffer.allocate(4);
    bb.putInt(0, 42);
    assertEquals(query.getRoutingKey(protocolVersion, codecRegistry), bb);
    session().execute(query);

    query = select().from(table).where(eq("ThEkEy", 42));
    assertEquals(query.getRoutingKey(protocolVersion, codecRegistry), bb);
    row = session().execute(query).one();
    assertEquals(row.getInt("theKey"), 42);
    assertEquals(row.getInt("a"), 1);
    assertEquals(row.getInt("b"), 2);
  }

  @Test(groups = "short")
  public void routingKeyColumnCaseSensitivityForQuotedIdentifiersTest() throws Exception {

    BuiltStatement query;
    TableMetadata table = cluster().getMetadata().getKeyspace(keyspace).getTable(TABLE_CASE_QUOTED);
    assertNotNull(table);
    ProtocolVersion protocolVersion =
        cluster().getConfiguration().getProtocolOptions().getProtocolVersion();
    CodecRegistry codecRegistry = CodecRegistry.DEFAULT_INSTANCE;

    query =
        insertInto(table)
            .values(
                new String[] {"\"theKey\"", "a", "b", "\"tHEkEY\""}, new Object[] {42, 1, 2, 3});
    ByteBuffer bb = ByteBuffer.allocate(4);
    bb.putInt(0, 42);
    assertEquals(query.getRoutingKey(protocolVersion, codecRegistry), bb);

    query =
        insertInto(table)
            .values(new String[] {"theKey", "a", "b", "\"tHEkEY\""}, new Object[] {42, 1, 2, 3});
    assertNull(query.getRoutingKey(protocolVersion, codecRegistry));

    query =
        insertInto(table)
            .values(new String[] {"theKey", "a", "b", "theKey"}, new Object[] {42, 1, 2, 3});
    assertNull(query.getRoutingKey(protocolVersion, codecRegistry));
  }

  @Test(groups = "short")
  public void intRoutingKeyTest() throws Exception {

    BuiltStatement query;
    TableMetadata table = cluster().getMetadata().getKeyspace(keyspace).getTable(TABLE_INT);
    assertNotNull(table);
    ProtocolVersion protocolVersion =
        cluster().getConfiguration().getProtocolOptions().getProtocolVersion();
    CodecRegistry codecRegistry = CodecRegistry.DEFAULT_INSTANCE;

    query = insertInto(table).values(new String[] {"k", "a", "b"}, new Object[] {42, 1, 2});
    ByteBuffer bb = ByteBuffer.allocate(4);
    bb.putInt(0, 42);
    assertEquals(query.getRoutingKey(protocolVersion, codecRegistry), bb);
    session().execute(query);

    query = select().from(table).where(eq("k", 42));
    assertEquals(query.getRoutingKey(protocolVersion, codecRegistry), bb);
    Row row = session().execute(query).one();
    assertEquals(row.getInt("k"), 42);
    assertEquals(row.getInt("a"), 1);
    assertEquals(row.getInt("b"), 2);
  }

  @Test(groups = "short")
  public void intRoutingBatchKeyTest() throws Exception {

    BuiltStatement query;
    TableMetadata table = cluster().getMetadata().getKeyspace(keyspace).getTable(TABLE_INT);
    assertNotNull(table);
    ProtocolVersion protocolVersion =
        cluster().getConfiguration().getProtocolOptions().getProtocolVersion();
    CodecRegistry codecRegistry = CodecRegistry.DEFAULT_INSTANCE;

    ByteBuffer bb = ByteBuffer.allocate(4);
    bb.putInt(0, 42);

    String batch_query;
    BuiltStatement batch;

    query = select().from(table).where(eq("k", 42));

    batch_query = "BEGIN BATCH ";
    batch_query += String.format("INSERT INTO %s.test_int (k,a) VALUES (42,1);", keyspace);
    batch_query += String.format("UPDATE %s.test_int USING TTL 400;", keyspace);
    batch_query += "APPLY BATCH;";
    batch =
        batch()
            .add(insertInto(table).values(new String[] {"k", "a"}, new Object[] {42, 1}))
            .add(update(table).using(ttl(400)));
    assertEquals(batch.getRoutingKey(protocolVersion, codecRegistry), bb);
    assertEquals(batch.toString(), batch_query);
    // TODO: rs = session().execute(batch); // Not guaranteed to be valid CQL

    batch_query = "BEGIN BATCH ";
    batch_query += String.format("SELECT * FROM %s.test_int WHERE k=42;", keyspace);
    batch_query += "APPLY BATCH;";
    batch = batch(query);
    assertEquals(batch.getRoutingKey(protocolVersion, codecRegistry), bb);
    assertEquals(batch.toString(), batch_query);
    // TODO: rs = session().execute(batch); // Not guaranteed to be valid CQL

    batch_query = "BEGIN BATCH ";
    batch_query += "SELECT * FROM foo WHERE k=42;";
    batch_query += "APPLY BATCH;";
    batch = batch().add(select().from("foo").where(eq("k", 42)));
    assertEquals(batch.getRoutingKey(protocolVersion, codecRegistry), null);
    assertEquals(batch.toString(), batch_query);
    // TODO: rs = session().execute(batch); // Not guaranteed to be valid CQL

    batch_query = "BEGIN BATCH USING TIMESTAMP 42 ";
    batch_query += "INSERT INTO foo.bar (a) VALUES (123);";
    batch_query += "APPLY BATCH;";
    batch = batch().using(timestamp(42)).add(insertInto("foo", "bar").value("a", 123));
    assertEquals(batch.getRoutingKey(protocolVersion, codecRegistry), null);
    assertEquals(batch.toString(), batch_query);
    // TODO: rs = session().execute(batch); // Not guaranteed to be valid CQL
  }
}
