/*
 * 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.DriverException;
import com.datastax.driver.core.policies.*;
import com.datastax.driver.core.utils.CassandraVersion;
import org.testng.annotations.Test;

import java.util.Iterator;
import java.util.concurrent.atomic.AtomicInteger;

import static org.assertj.core.api.Assertions.assertThat;

public class StatementWrapperTest extends CCMTestsSupport {

    private static final String INSERT_QUERY = "insert into test (k, v) values (?, ?)";
    private static final String SELECT_QUERY = "select * from test where k = ?";

    private static final String INSERT_MULTIPAGE_QUERY = "insert into test_multipage (k, v) values (?, ?)";
    private static final String SELECT_MULTIPAGE_QUERY = "select * from test_multipage where k = ?";

    CustomLoadBalancingPolicy loadBalancingPolicy = new CustomLoadBalancingPolicy();
    CustomSpeculativeExecutionPolicy speculativeExecutionPolicy = new CustomSpeculativeExecutionPolicy();
    CustomRetryPolicy retryPolicy = new CustomRetryPolicy();

    @Override
    public void onTestContextInitialized() {
        execute("create table test (k text primary key, v int)");
        execute("create table test_multipage (k text, v int, primary key (k, v))");
    }

    @Override
    public Cluster.Builder createClusterBuilder() {
        return Cluster.builder()
                .withLoadBalancingPolicy(loadBalancingPolicy)
                .withSpeculativeExecutionPolicy(speculativeExecutionPolicy)
                .withRetryPolicy(retryPolicy);
    }

    @Test(groups = "short")
    public void should_pass_wrapped_statement_to_load_balancing_policy() {
        loadBalancingPolicy.customStatementsHandled.set(0);

        SimpleStatement s = new SimpleStatement("select * from system.local");
        session().execute(s);
        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(0);

        session().execute(new CustomStatement(s));
        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);
    }

    @Test(groups = "short")
    @CassandraVersion(value = "2.0.0")
    public void should_reuse_wrapped_simple_statement_for_multipage_query() {
        loadBalancingPolicy.customStatementsHandled.set(0);

        for (int v = 1; v <= 100; v++)
            session().execute(new SimpleStatement(INSERT_MULTIPAGE_QUERY, "key_simple_multipage", v));

        SimpleStatement s = new SimpleStatement(SELECT_MULTIPAGE_QUERY, "key_simple_multipage");
        s.setFetchSize(1);

        CustomStatement customStatement = new CustomStatement(s);

        ResultSet rs = session().execute(customStatement);
        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

        Iterator<Row> it = rs.iterator();

        assertThat(it.hasNext()).isTrue();
        it.next();
        assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

        assertThat(it.hasNext()).isTrue();
        it.next();
        assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);

        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(2);
    }

    @Test(groups = "short")
    @CassandraVersion(value = "2.0.0")
    public void should_reuse_wrapped_bound_statement_for_multipage_query() {
        loadBalancingPolicy.customStatementsHandled.set(0);

        for (int v = 1; v <= 100; v++)
            session().execute(new SimpleStatement(INSERT_MULTIPAGE_QUERY, "key_prepared_multipage", v));

        PreparedStatement ps = session().prepare(SELECT_MULTIPAGE_QUERY);
        BoundStatement bs = ps.bind("key_prepared_multipage");
        bs.setFetchSize(1);

        CustomStatement customStatement = new CustomStatement(bs);

        ResultSet rs = session().execute(customStatement);
        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

        Iterator<Row> it = rs.iterator();

        assertThat(it.hasNext()).isTrue();
        it.next();
        assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);
        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(1);

        assertThat(it.hasNext()).isTrue();
        it.next();
        assertThat(rs.getExecutionInfo().getStatement()).isEqualTo(customStatement);

        assertThat(loadBalancingPolicy.customStatementsHandled.get()).isEqualTo(2);
    }

    @Test(groups = "short")
    public void should_pass_wrapped_statement_to_speculative_execution_policy() {
        speculativeExecutionPolicy.customStatementsHandled.set(0);

        SimpleStatement s = new SimpleStatement("select * from system.local");
        session().execute(s);
        assertThat(speculativeExecutionPolicy.customStatementsHandled.get()).isEqualTo(0);

        session().execute(new CustomStatement(s));
        assertThat(speculativeExecutionPolicy.customStatementsHandled.get()).isEqualTo(1);
    }

    @Test(groups = "short")
    public void should_pass_wrapped_statement_to_retry_policy() {
        retryPolicy.customStatementsHandled.set(0);

        // Set CL TWO with only one node, so the statement will always cause UNAVAILABLE,
        // which our custom policy ignores.
        Statement s = new SimpleStatement("select * from system.local")
                .setConsistencyLevel(ConsistencyLevel.TWO);

        session().execute(s);
        assertThat(retryPolicy.customStatementsHandled.get()).isEqualTo(0);

        session().execute(new CustomStatement(s));
        assertThat(retryPolicy.customStatementsHandled.get()).isEqualTo(1);
    }

    @CassandraVersion("2.0.0")
    @Test(groups = "short")
    public void should_execute_wrapped_simple_statement() {
        session().execute(new CustomStatement(new SimpleStatement(INSERT_QUERY, "key_simple", 1)));

        ResultSet rs = session().execute(new CustomStatement(new SimpleStatement(SELECT_QUERY, "key_simple")));
        assertThat(rs.one().getInt("v")).isEqualTo(1);
    }

    @Test(groups = "short")
    public void should_execute_wrapped_bound_statement() {
        PreparedStatement preparedStatement = session().prepare(new SimpleStatement(INSERT_QUERY));
        session().execute(new CustomStatement(preparedStatement.bind("key_bound", 1)));

        preparedStatement = session().prepare(new SimpleStatement(SELECT_QUERY));
        ResultSet rs = session().execute(new CustomStatement(preparedStatement.bind("key_bound")));
        assertThat(rs.one().getInt("v")).isEqualTo(1);
    }

    @CassandraVersion("2.0.0")
    @Test(groups = "short")
    public void should_execute_wrapped_batch_statement() {
        BatchStatement batchStatement = new BatchStatement();
        batchStatement.add(new SimpleStatement(INSERT_QUERY, "key_batch", 1));

        session().execute(new CustomStatement(batchStatement));

        ResultSet rs = session().execute(SELECT_QUERY, "key_batch");
        assertThat(rs.one().getInt("v")).isEqualTo(1);
    }

    @CassandraVersion("2.0.0")
    @Test(groups = "short")
    public void should_add_wrapped_batch_statement_to_batch_statement() {
        BatchStatement batchStatementForWrapping = new BatchStatement();
        batchStatementForWrapping.add(new SimpleStatement(INSERT_QUERY, "key1", 1));

        BatchStatement batchStatement = new BatchStatement();
        batchStatement.add(new CustomStatement(new SimpleStatement(INSERT_QUERY, "key2", 2)));
        batchStatement.add(new CustomStatement(batchStatementForWrapping));

        session().execute(batchStatement);

        ResultSet rs = session().execute(SELECT_QUERY, "key1");
        assertThat(rs.one().getInt("v")).isEqualTo(1);

        rs = session().execute(SELECT_QUERY, "key2");
        assertThat(rs.one().getInt("v")).isEqualTo(2);
    }

    /**
     * A custom wrapper that's just used to mark statements.
     */
    static class CustomStatement extends StatementWrapper {
        protected CustomStatement(Statement wrapped) {
            super(wrapped);
        }
    }

    /**
     * A load balancing policy that counts how many times it has seen the custom wrapper
     */
    static class CustomLoadBalancingPolicy extends DelegatingLoadBalancingPolicy {
        final AtomicInteger customStatementsHandled = new AtomicInteger();

        public CustomLoadBalancingPolicy() {
            super(new RoundRobinPolicy());
        }

        @Override
        public Iterator<Host> newQueryPlan(String loggedKeyspace, Statement statement) {
            if (statement instanceof CustomStatement)
                customStatementsHandled.incrementAndGet();
            return super.newQueryPlan(loggedKeyspace, statement);
        }
    }

    /**
     * A speculative execution policy that counts how many times it has seen the custom wrapper
     */
    static class CustomSpeculativeExecutionPolicy extends DelegatingSpeculativeExecutionPolicy {
        final AtomicInteger customStatementsHandled = new AtomicInteger();

        protected CustomSpeculativeExecutionPolicy() {
            super(NoSpeculativeExecutionPolicy.INSTANCE);
        }

        @Override
        public SpeculativeExecutionPlan newPlan(String loggedKeyspace, Statement statement) {
            if (statement instanceof CustomStatement)
                customStatementsHandled.incrementAndGet();
            return super.newPlan(loggedKeyspace, statement);
        }
    }

    /**
     * A retry policy that counts how many times it has seen the custom wrapper for UNAVAILABLE errors.
     */
    static class CustomRetryPolicy implements RetryPolicy {
        final AtomicInteger customStatementsHandled = new AtomicInteger();

        @Override
        public RetryDecision onUnavailable(Statement statement, ConsistencyLevel cl, int requiredReplica, int aliveReplica, int nbRetry) {
            if (statement instanceof CustomStatement)
                customStatementsHandled.incrementAndGet();
            return RetryDecision.ignore();
        }

        @Override
        public RetryDecision onReadTimeout(Statement statement, ConsistencyLevel cl, int requiredResponses, int receivedResponses, boolean dataRetrieved, int nbRetry) {
            return RetryDecision.rethrow();
        }

        @Override
        public RetryDecision onWriteTimeout(Statement statement, ConsistencyLevel cl, WriteType writeType, int requiredAcks, int receivedAcks, int nbRetry) {
            return RetryDecision.rethrow();
        }

        @Override
        public RetryDecision onRequestError(Statement statement, ConsistencyLevel cl, DriverException e, int nbRetry) {
            return RetryDecision.tryNextHost(cl);
        }

        @Override
        public void init(Cluster cluster) {
            // nothing to do
        }

        @Override
        public void close() {
            // nothing to do
        }
    }
}
