/*
 * 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.querybuilder.QueryBuilder.bindMarker;
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 org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Fail.fail;

import com.datastax.driver.core.exceptions.FrameTooLongException;
import com.datastax.driver.core.querybuilder.Insert;
import com.datastax.driver.core.schemabuilder.Create;
import com.datastax.driver.core.schemabuilder.SchemaBuilder;
import com.datastax.driver.core.utils.CassandraVersion;
import java.nio.ByteBuffer;
import java.util.Collection;
import java.util.Random;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

@CCMConfig(numberOfNodes = 2)
public class FrameLengthTest extends CCMTestsSupport {

  Logger logger = LoggerFactory.getLogger(FrameLengthTest.class);

  private static final String tableName = "blob_table";
  private static final int colCount = 256;
  private static final int rowsPerPartitionCount = 4;
  private static final int partitionCount = 1;
  private static final int bytesPerCol = 1024;

  @BeforeClass(groups = {"isolated"})
  public void beforeTestClass() throws Exception {
    // Set max frame size to 1MB to make it easier to manifest frame length error.
    System.setProperty("com.datastax.driver.NATIVE_TRANSPORT_MAX_FRAME_SIZE_IN_MB", "1");
    super.beforeTestClass();
  }

  @Override
  public void onTestContextInitialized() {
    logger.info("Creating table {} with {} {}-byte blob columns", tableName, colCount, bytesPerCol);
    Random random = new Random();
    // Create table
    Create create =
        SchemaBuilder.createTable(tableName)
            .addPartitionKey("k", DataType.cint())
            .addClusteringColumn("c", DataType.cint());
    for (int i = 0; i < colCount; i++) {
      create.addColumn("col" + i, DataType.blob());
    }
    execute(create.getQueryString());

    // build prepared statement.
    Insert insert = insertInto(tableName).value("k", bindMarker()).value("c", bindMarker());
    for (int i = 0; i < colCount; i++) {
      insert = insert.value("col" + i, bindMarker());
    }

    PreparedStatement prepared = session().prepare(insert);

    // Insert rows.
    logger.info("Inserting data for {} partitions.", partitionCount);
    for (int i = 0; i < partitionCount; i++) {
      logger.info("Inserting {} rows in partition {}", rowsPerPartitionCount, i);
      for (int r = 0; r < rowsPerPartitionCount; r++) {
        BoundStatement stmt = prepared.bind();
        stmt.setInt("k", i);
        stmt.setInt("c", r);
        for (int c = 0; c < colCount; c++) {
          byte[] b = new byte[bytesPerCol];
          random.nextBytes(b);
          ByteBuffer in = ByteBuffer.wrap(b);
          stmt.setBytes("col" + c, in);
        }
        session().execute(stmt);
      }
    }
    logger.info("Done loading {}", tableName);
  }

  /**
   * Validates that if a frame is received that exceeds NATIVE_TRANSPORT_MAX_FRAME_SIZE_IN_MB that
   * the driver is able to recover, not lose host connectivity and make further queries. It
   * configures NATIVE_TRANSPORT_MAX_FRAME_SIZE_IN_MB to 1 MB to make the error easier to reproduce.
   *
   * @jira_ticket JAVA-1292
   * @jira_ticket JAVA-1293
   * @test_category connection
   */
  @Test(groups = "isolated")
  public void should_throw_exception_when_frame_exceeds_configured_max() {
    skipTestWithCassandraVersionOrHigher("4.0.0", "frame-size exceeding with default-protocol");
    run_frame_size_exceeding_queries(session());
  }

  /**
   * With cassandra 4.0.0+, V5 protocol is default which breaks requests into segments. Force V4
   * protocol to allow us to test frame-size limitation code.
   */
  @CassandraVersion("4.0.0")
  @Test(groups = "isolated")
  public void should_throw_exception_when_frame_exceeds_configured_max_v4_protocol_cassandra4() {
    Cluster cluster =
        register(createClusterBuilder().withProtocolVersion(ProtocolVersion.V4).build());
    Session session = register(cluster.connect());
    useKeyspace(session, keyspace);

    run_frame_size_exceeding_queries(session);
  }

  private void run_frame_size_exceeding_queries(Session session) {
    try {
      session.execute(select().from(tableName).where(eq("k", 0)));
      fail("Exception expected");
    } catch (FrameTooLongException ftle) {
      // Expected.
    }

    // Both hosts should remain up.
    Collection<Host> hosts = session.getState().getConnectedHosts();
    assertThat(hosts).hasSize(2).extractingResultOf("isUp").containsOnly(true);

    // Should be able to make a query that is less than the max frame size.
    // Execute multiple time to exercise all hosts.
    for (int i = 0; i < 10; i++) {
      ResultSet result =
          session.execute(select().from(tableName).where(eq("k", 0)).and(eq("c", 0)));
      assertThat(result.getAvailableWithoutFetching()).isEqualTo(1);
    }
  }
}
