/*
 * 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.ProtocolVersion.V3;
import static com.datastax.driver.core.querybuilder.QueryBuilder.eq;
import static com.datastax.driver.core.querybuilder.QueryBuilder.select;
import static org.apache.log4j.Level.TRACE;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Fail.fail;

import com.datastax.driver.core.exceptions.UnsupportedFeatureException;
import com.datastax.driver.core.utils.CassandraVersion;
import com.google.common.collect.ImmutableMap;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import org.apache.log4j.Logger;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

@CassandraVersion("2.2.0")
@CCMConfig(
    jvmArgs =
        "-Dcassandra.custom_query_handler_class=org.apache.cassandra.cql3.CustomPayloadMirroringQueryHandler")
public class CustomPayloadTest extends CCMTestsSupport {

  private Map<String, ByteBuffer> payload1;

  private Map<String, ByteBuffer> payload2;

  @BeforeMethod(groups = {"short", "unit"})
  public void initPayloads() {
    payload1 =
        ImmutableMap.of(
            "k1", ByteBuffer.wrap(new byte[] {1, 2, 3}),
            "k2", ByteBuffer.wrap(new byte[] {4, 5, 6}));
    payload2 =
        ImmutableMap.of(
            "k2", ByteBuffer.wrap(new byte[] {1, 2}),
            "k3", ByteBuffer.wrap(new byte[] {3, 4}));
  }

  // execute

  @Test(groups = "short")
  public void should_echo_custom_payload_when_executing_statement() throws Exception {
    Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1);
    statement.setOutgoingPayload(payload1);
    ResultSet rows = session().execute(statement);
    Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload1);
  }

  @Test(groups = "short")
  public void should_echo_custom_payload_when_executing_batch_statement() throws Exception {
    Statement statement =
        new BatchStatement().add(new SimpleStatement("INSERT INTO t1 (c1, c2) values (1, 'foo')"));
    statement.setOutgoingPayload(payload1);
    ResultSet rows = session().execute(statement);
    Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload1);
  }

  @Test(groups = "short")
  public void should_echo_custom_payload_when_building_statement() throws Exception {
    Statement statement = select("c2").from("t1").where(eq("c1", 1)).setOutgoingPayload(payload1);
    ResultSet rows = session().execute(statement);
    Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload1);
  }

  // prepare

  /** Ensures that an incoming payload is propagated from prepared to bound statements. */
  @Test(groups = "short")
  public void should_propagate_incoming_payload_to_bound_statement() throws Exception {
    RegularStatement statement = new SimpleStatement("SELECT c2 as col1 FROM t1 where c1 = ?");
    statement.setOutgoingPayload(payload1);
    PreparedStatement ps = session().prepare(statement);
    // Prepared statement should inherit outgoing payload
    assertThat(ps.getOutgoingPayload()).isEqualTo(payload1);
    // Prepared statement should receive incoming payload
    assertThat(ps.getIncomingPayload()).isEqualTo(payload1);
    ps.setOutgoingPayload(null); // unset outgoing payload
    // bound statement should inherit from prepared statement's incoming payload
    BoundStatement bs = ps.bind(1);
    ResultSet rows = session().execute(bs);
    Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload1);
    bs = ps.bind();
    bs.setInt(0, 1);
    rows = session().execute(bs);
    actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload1);
  }

  /**
   * Ensures that an incoming payload is overridden by an explicitly set outgoing payload when
   * propagated to bound statements.
   */
  @Test(groups = "short")
  public void
      should_override_incoming_payload_when_outgoing_payload_explicitly_set_on_preparing_statement()
          throws Exception {
    RegularStatement statement = new SimpleStatement("SELECT c2 as col2 FROM t1 where c1 = ?");
    statement.setOutgoingPayload(payload1);
    PreparedStatement ps = session().prepare(statement);
    // Prepared statement should inherit outgoing payload
    assertThat(ps.getOutgoingPayload()).isEqualTo(payload1);
    // Prepared statement should receive incoming payload
    assertThat(ps.getIncomingPayload()).isEqualTo(payload1);
    ps.setOutgoingPayload(payload2); // override outgoing payload
    // bound statement should inherit from prepared statement's outgoing payload
    BoundStatement bs = ps.bind(1);
    ResultSet rows = session().execute(bs);
    Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload2);
    bs = ps.bind();
    bs.setInt(0, 1);
    rows = session().execute(bs);
    actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload2);
  }

  /**
   * Ensures that payloads can still be set individually on bound statements if the prepared
   * statement does not have a default payload.
   */
  @Test(groups = "short")
  public void should_not_set_any_payload_on_bound_statement() throws Exception {
    RegularStatement statement = new SimpleStatement("SELECT c2 as col3 FROM t1 where c1 = ?");
    PreparedStatement ps = session().prepare(statement);
    assertThat(ps.getOutgoingPayload()).isNull();
    assertThat(ps.getIncomingPayload()).isNull();
    // bound statement should not have outgoing payload
    BoundStatement bs = ps.bind(1);
    assertThat(bs.getOutgoingPayload()).isNull();
    // explicitly set a payload for this boudn statement only
    bs.setOutgoingPayload(payload1);
    ResultSet rows = session().execute(bs);
    Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload1);
    // a second bound statement should not have any payload
    bs = ps.bind();
    assertThat(bs.getOutgoingPayload()).isNull();
    bs.setInt(0, 1);
    rows = session().execute(bs);
    actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isNull();
  }

  // pagination

  /** Ensures that a custom payload is propagated throughout pages. */
  @Test(groups = "short")
  public void should_echo_custom_payload_when_paginating() throws Exception {
    session().execute("INSERT INTO t1 (c1, c2) VALUES (1, 'a')");
    session().execute("INSERT INTO t1 (c1, c2) VALUES (1, 'b')");
    Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = 1");
    statement.setFetchSize(1);
    statement.setOutgoingPayload(payload1);
    ResultSet rows = session().execute(statement);
    rows.all();
    assertThat(rows.getAllExecutionInfo()).extracting("incomingPayload").containsOnly(payload1);
  }

  // TODO retries, spec execs

  // edge cases

  @Test(groups = "short")
  public void should_encode_null_values() throws Exception {
    Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>();
    payload.put("k1", Statement.NULL_PAYLOAD_VALUE);
    Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1);
    statement.setOutgoingPayload(payload);
    ResultSet rows = session().execute(statement);
    Map<String, ByteBuffer> actual = rows.getExecutionInfo().getIncomingPayload();
    assertThat(actual).isEqualTo(payload);
  }

  @Test(groups = "unit", expectedExceptions = NullPointerException.class)
  public void should_throw_npe_when_null_key_on_regular_statement() throws Exception {
    Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>();
    payload.put(null, ByteBuffer.wrap(new byte[] {1}));
    new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1).setOutgoingPayload(payload);
  }

  @Test(groups = "unit", expectedExceptions = NullPointerException.class)
  public void should_throw_npe_when_null_value_on_regular_statement() throws Exception {
    Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>();
    payload.put("k1", null);
    new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1).setOutgoingPayload(payload);
  }

  @Test(groups = "short", expectedExceptions = NullPointerException.class)
  public void should_throw_npe_when_null_key_on_prepared_statement() throws Exception {
    Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>();
    payload.put(null, ByteBuffer.wrap(new byte[] {1}));
    session()
        .prepare(new SimpleStatement("SELECT c2 FROM t1 where c1 = 1"))
        .setOutgoingPayload(payload);
  }

  @Test(groups = "short", expectedExceptions = NullPointerException.class)
  public void should_throw_npe_when_null_value_on_prepared_statement() throws Exception {
    Map<String, ByteBuffer> payload = new HashMap<String, ByteBuffer>();
    payload.put("k1", null);
    session()
        .prepare(new SimpleStatement("SELECT c2 FROM t1 where c1 = 2"))
        .setOutgoingPayload(payload);
  }

  @Test(groups = "short")
  public void should_throw_ufe_when_protocol_version_lesser_than_4() throws Exception {
    try {
      Cluster v3cluster =
          register(
                  Cluster.builder()
                      .addContactPoints(getContactPoints())
                      .withPort(ccm().getBinaryPort())
                      .withProtocolVersion(V3)
                      .build())
              .init();
      Session v3session = v3cluster.connect();
      Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1);
      statement.setOutgoingPayload(payload1);
      v3session.execute(statement);
      fail("Should not send custom payloads with protocol V3");
    } catch (UnsupportedFeatureException e) {
      assertThat(e.getMessage())
          .isEqualTo(
              "Unsupported feature with the native protocol V3 (which is currently in use): Custom payloads are only supported since native protocol V4");
    }
  }

  // log messages

  /**
   * Ensures that when debugging custom payloads, the driver will print appropriate log messages.
   */
  @Test(groups = "short")
  public void should_print_log_message_when_level_trace() throws Exception {
    Logger logger = Logger.getLogger(Message.logger.getName());
    MemoryAppender appender = new MemoryAppender();
    try {
      logger.setLevel(TRACE);
      logger.addAppender(appender);
      Statement statement = new SimpleStatement("SELECT c2 FROM t1 where c1 = ?", 1);
      statement.setOutgoingPayload(payload1);
      session().execute(statement);
      String logs = appender.waitAndGet(10000);
      assertThat(logs)
          .contains("Sending payload: {k1:0x010203, k2:0x040506} (24 bytes total)")
          .contains("Received payload: {k1:0x010203, k2:0x040506} (24 bytes total)");
    } finally {
      logger.setLevel(null);
      logger.removeAppender(appender);
    }
  }

  @Override
  public void onTestContextInitialized() {
    execute("CREATE TABLE t1 (c1 int, c2 text,  PRIMARY KEY (c1, c2))");
  }
}
