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

import com.datastax.driver.core.exceptions.UnsupportedFeatureException;
import com.datastax.driver.core.utils.CassandraVersion;
import com.google.common.collect.ImmutableMap;
import org.apache.log4j.Logger;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;

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;

@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))");
    }

}
