/*
 * 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 org.apache.cassandra.distributed.test.guardrails;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Stream;

import com.datastax.driver.core.Session;
import com.datastax.driver.core.exceptions.InvalidQueryException;
import com.google.common.util.concurrent.Uninterruptibles;

import org.awaitility.Awaitility;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

import org.apache.cassandra.cql3.QueryProcessor;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.Keyspace;
import org.apache.cassandra.db.guardrails.Guardrails;
import org.apache.cassandra.distributed.Cluster;
import org.apache.cassandra.distributed.api.ConsistencyLevel;
import org.apache.cassandra.distributed.api.Feature;
import org.apache.cassandra.distributed.api.IInvokableInstance;
import org.apache.cassandra.distributed.api.IIsolatedExecutor;
import org.apache.cassandra.exceptions.QueryReferencesTooManyIndexesAbortException;
import org.apache.cassandra.exceptions.ReadFailureException;
import org.apache.cassandra.index.Index;
import org.apache.cassandra.index.SecondaryIndexManager;
import org.apache.cassandra.service.ClientWarn;
import org.apache.cassandra.service.QueryState;
import org.apache.cassandra.service.reads.thresholds.CoordinatorWarnings;

import static com.google.common.collect.Iterables.getOnlyElement;
import static org.apache.cassandra.db.ConsistencyLevel.ALL;
import static org.apache.cassandra.service.reads.thresholds.WarningsSnapshot.tooManyIndexesReadWarnMessage;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

@SuppressWarnings("Convert2MethodRef")
public class GuardrailNonPartitionRestrictedQueryTest extends GuardrailTester
{
    private static Cluster cluster;
    private static com.datastax.driver.core.Cluster driverCluster;
    private Session driverSession;

    @BeforeClass
    public static void setupCluster() throws IOException
    {
        cluster = Cluster.build(3)
                         .withConfig(c -> c.with(Feature.GOSSIP, Feature.NATIVE_PROTOCOL)
                                           .set("read_thresholds_enabled", "true")
                                           .set("authenticator", "PasswordAuthenticator"))
                         .withDataDirCount(1)
                         .start();

        driverCluster = buildDriverCluster(cluster);
    }

    @AfterClass
    public static void teardownCluster()
    {
        if (driverCluster != null)
            driverCluster.close();

        if (cluster != null)
            cluster.close();
    }

    @Before
    public void beforeTest()
    {
        super.beforeTest();
        cluster.schemaChange("DROP KEYSPACE IF EXISTS " + KEYSPACE);
        init(cluster);
        driverSession = driverCluster.connect();
    }

    @After
    public void afterTest()
    {
        if (driverSession != null)
            driverSession.close();
    }

    @Override
    protected Cluster getCluster()
    {
        return cluster;
    }

    @Override
    protected Session getSession()
    {
        return driverSession;
    }

    @Test
    public void testGuardrailForLegacy2i()
    {
        prepareSchema(false);
        testGuardrailInternal();
    }

    @Test
    public void testGuardrailForSAI()
    {
        prepareSchema(true);
        testGuardrailInternal();
    }

    @Test
    public void testSAIWarnThreshold()
    {
        prepareSchema(true);
        setThresholds(2, 5);

        // flushing just 1 table will not trigger any threshold
        long valueToQuery = createSSTables(1);
        assertThat(executeSelect(valueToQuery, false)).isNull();

        assertWarnAborts(0, 0);

        // create 3 more SSTables on each node, this will trigger warn threshold (3 > 2 but < 5)
        valueToQuery = createSSTables(3);
        String expectedMessage = tooManyIndexesReadWarnMessage(cluster.size(),
                                                               3,
                                                               String.format("SELECT * FROM %s.%s WHERE v1 = %s ALLOW FILTERING",
                                                                             KEYSPACE, tableName, valueToQuery));
        assertThat(getOnlyElement(executeSelect(valueToQuery, false))).contains(expectedMessage);

        assertWarnAborts(1, 0);

        // we effectively stop reacting on warnings
        setThresholds(32, -1);
        assertThat(executeSelect(valueToQuery, false)).isNull();

        assertWarnAborts(1, 0);

        // we compacted SSTables of node 1 but not on 2 and 3,
        // we set tresholds to all nodes hence we violated thresholds on node 2 and 3
        compact(1);
        setThresholds(2, 5);

        // notice we expect warnings from 2 nodes
        expectedMessage = tooManyIndexesReadWarnMessage(cluster.size() - 1,
                                                        3,
                                                        String.format("SELECT * FROM %s.%s WHERE v1 = %s ALLOW FILTERING",
                                                                      KEYSPACE, tableName, valueToQuery));

        assertThat(getOnlyElement(executeSelect(valueToQuery, false))).contains(expectedMessage);

        assertWarnAborts(2, 0);

        // disable warn thresholds on nodes 2 and 3
        setThresholds(32, -1, 2, 3);
        // we will not warn because nodes 2 and 3 are disabled and 1 does not violate warn threshold
        assertThat(executeSelect(valueToQuery, false)).isNull();

        assertWarnAborts(2, 0);

        // set thresholds back on nodes 2 and 3
        setThresholds(2, 5, 2, 3);
        assertThat(getOnlyElement(executeSelect(valueToQuery, false))).contains(expectedMessage);

        assertWarnAborts(3, 0);

        // we compacted SSTables on the third node as well, so we will not violate any thresholds
        compact(2);
        compact(3);

        assertThat(executeSelect(valueToQuery, false)).isNull();

        // metrics stay same as they were
        assertWarnAborts(3, 0);
    }

    @Test
    public void testSAIFailThreshold()
    {
        prepareSchema(true);
        setThresholds(2, 5);

        // create 6 SSTables on each node, this will trigger fail threshold (6 > 5)
        long valueToQuery = createSSTables(6);

        String expectedMessage = String.format("referenced %s SSTable indexes for a query without restrictions on partition key " +
                                               "and aborted the query SELECT * FROM %s.%s WHERE v1 = %s ALLOW FILTERING",
                                               6, KEYSPACE, tableName, valueToQuery);

        assertThat(getOnlyElement(executeSelect(valueToQuery, true))).contains(expectedMessage);

        assertWarnAborts(0, 1);

        // we effectively stop reacting on failures
        setThresholds(32, -1);
        assertThat(executeSelect(valueToQuery, false)).isNull();

        // we compacted SSTables of node 1 but not on 2 and 3,
        // we set tresholds to all nodes hence we violated thresholds on node 2 and 3
        compact(1);
        setThresholds(2, 5);
        assertThat(getOnlyElement(executeSelect(valueToQuery, true))).contains(expectedMessage);

        assertWarnAborts(0, 2);

        // disable fail thresholds on nodes 2 and 3
        setThresholds(32, -1, 2, 3);
        // we will not fail because nodes 2 and 3 are disabled and 1 does not violate fail threshold
        assertThat(executeSelect(valueToQuery, false)).isNull();

        // set thresholds back on nodes 2 and 3
        setThresholds(2, 5, 2, 3);
        assertThat(getOnlyElement(executeSelect(valueToQuery, true))).contains(expectedMessage);

        assertWarnAborts(0, 3);

        // we compacted SSTables on the second and the third node as well, so we will not violate any thresholds
        compact(2);
        compact(3);

        assertThat(executeSelect(valueToQuery, false)).isNull();

        // metrics stay same as they were
        assertWarnAborts(0, 3);
    }

    private void testGuardrailInternal()
    {
        enableGuardrail();

        Awaitility.await()
                  .pollDelay(5, TimeUnit.SECONDS)
                  .atMost(1, TimeUnit.MINUTES)
                  .pollInterval(5, TimeUnit.SECONDS)
                  .until(() -> {
                      try
                      {
                          assertThat(executeViaDriver(String.format("SELECT * from %s.%s WHERE k = 0 AND v1 = 0", KEYSPACE, tableName))).isEmpty();
                          assertThat(executeViaDriver(String.format("SELECT * from %s.%s WHERE v1 = 0", KEYSPACE, tableName))).isEmpty();

                          return true;
                      }
                      catch (ReadFailureException ex)
                      {
                          return false;
                      }
                  });

        disableGuardrail();

        try
        {
            executeViaDriver(String.format("SELECT * from %s.%s WHERE v1 = 0", KEYSPACE, tableName));
            fail("selection on non-partition restricted queries should be forbidden");
        }
        catch (InvalidQueryException e)
        {
            assertThat(e.getMessage()).contains(Guardrails.nonPartitionRestrictedIndexQueryEnabled.reason);
        }

        // even we disabled guardrail, if we query by primary key and clustering column, it passes
        assertThat(executeViaDriver(String.format("SELECT * from %s.%s WHERE k = 0 AND v1 = 0", KEYSPACE, tableName))).isEmpty();
        assertThat(executeViaDriver(String.format("SELECT * from %s.%s WHERE k = 0 AND c = 0", KEYSPACE, tableName))).isEmpty();
        assertThat(executeViaDriver(String.format("SELECT * from %s.%s WHERE k = 0", KEYSPACE, tableName))).isEmpty();

        // enable it back and do non-partition key queries
        enableGuardrail();

        assertThat(executeViaDriver(String.format("SELECT * from %s.%s WHERE v1 = 0", KEYSPACE, tableName))).isEmpty();
    }

    private void assertWarnAborts(int warns, int aborts)
    {
        assertThat(totalWarnings()).as("warnings").isEqualTo(warns);
        assertThat(totalAborts()).as("aborts").isEqualTo(aborts);
    }

    private List<String> executeSelect(long valueToQuery, boolean expectToFail)
    {
        return cluster.get(1).applyOnInstance((IIsolatedExecutor.SerializableTriFunction<String, String, Long, List<String>>) (keyspace, table, v1) -> {
            Awaitility.await()
                      .pollInterval(5, TimeUnit.SECONDS)
                      .atMost(1, TimeUnit.MINUTES)
                      .until(() -> {
                          ColumnFamilyStore cs = Keyspace.open(keyspace).getColumnFamilyStore(table);
                          if (cs == null)
                              return false;

                          SecondaryIndexManager indexManager = cs.indexManager;
                          if (indexManager == null)
                              return false;

                          Index v1Idx = indexManager.getIndexByName("v1_idx");
                          Index v2Idx = indexManager.getIndexByName("v2_idx");

                          if (v1Idx == null || v2Idx == null)
                              return false;

                          return indexManager.isIndexQueryable(v1Idx) && indexManager.isIndexQueryable(v2Idx);
                      });

            ClientWarn.instance.captureWarnings();
            CoordinatorWarnings.init();

            String query = String.format("SELECT * from %s.%s WHERE v1 = %s", keyspace, table, v1);

            try
            {
                QueryProcessor.execute(query, ALL, QueryState.forInternalCalls());
                if (expectToFail)
                    fail("expected to fail");
            }
            catch (QueryReferencesTooManyIndexesAbortException e)
            {
                if (!expectToFail)
                    fail("did not expect to fail");

                assertTrue(e.nodes >= 1 && e.nodes <= 3);
            }

            CoordinatorWarnings.done();
            CoordinatorWarnings.reset();
            return ClientWarn.instance.getWarnings();
        }, KEYSPACE, tableName, valueToQuery);
    }

    private boolean indexesQueryable()
    {
        for (int i = 1; i < cluster.size() + 1; i++)
        {
            boolean indexesQueryable = cluster.get(i).applyOnInstance((IIsolatedExecutor.SerializableBiFunction<String, String, Boolean>) (ks, tb) -> {
                for (ColumnFamilyStore cs : Keyspace.open(ks).getColumnFamilyStores())
                {
                    if (!cs.name.equals(tb))
                        continue;

                    SecondaryIndexManager indexManager = cs.indexManager;

                    if (indexManager == null)
                        return false;

                    Index v1Idx = indexManager.getIndexByName("v1_idx");
                    Index v2Idx = indexManager.getIndexByName("v2_idx");

                    if (v1Idx == null || v2Idx == null)
                        return false;

                    return indexManager.isIndexQueryable(v1Idx) && indexManager.isIndexQueryable(v2Idx);
                }

                return false;
            }, KEYSPACE, tableName);

            if (!indexesQueryable)
                return false;
        }

        return true;
    }

    private void prepareSchema(boolean sai)
    {
        schemaChange("CREATE TABLE IF NOT EXISTS %s (k bigint, c bigint, v1 bigint, v2 bigint, PRIMARY KEY (k, c));");

        if (sai)
        {
            schemaChange("CREATE CUSTOM INDEX IF NOT EXISTS v1_idx ON %s (v1) USING 'StorageAttachedIndex'");
            schemaChange("CREATE CUSTOM INDEX IF NOT EXISTS v2_idx ON %s (v2) USING 'StorageAttachedIndex'");
        }
        else
        {
            schemaChange("CREATE INDEX IF NOT EXISTS v1_idx ON %s (v1)");
            schemaChange("CREATE INDEX IF NOT EXISTS v2_idx ON %s (v2)");
        }

        Awaitility.await().atMost(1, TimeUnit.MINUTES).until(this::indexesQueryable);

        for (int i = 1; i < cluster.size() + 1; i++)
        {
            cluster.get(i).acceptsOnInstance((IIsolatedExecutor.SerializableBiConsumer<String, String>) (ks, tb) -> {
                for (ColumnFamilyStore cs : Keyspace.open(ks).getColumnFamilyStores())
                {
                    if (cs.name.equals(tb))
                    {
                        cs.disableAutoCompaction();
                        break;
                    }
                }
            }).accept(KEYSPACE, tableName);
        }
    }

    private void compact(int node)
    {
        cluster.get(node).acceptsOnInstance((IIsolatedExecutor.SerializableBiConsumer<String, String>) (ks, tb) -> {
            for (ColumnFamilyStore cs : Keyspace.open(ks).getColumnFamilyStores())
            {
                if (cs.name.equals(tb))
                {
                    cs.forceMajorCompaction();
                    break;
                }
            }
        }).accept(KEYSPACE, tableName);
    }

    private long createSSTables(int numberOfSSTables, int... nodesToFlush)
    {
        Uninterruptibles.sleepUninterruptibly(1, TimeUnit.SECONDS);

        long k = System.currentTimeMillis();
        long value = k + 10;
        for (int i = 0; i < numberOfSSTables; i++)
        {
            // this will be replicated to each node
            cluster.coordinator(1)
                   .execute(format("INSERT INTO %s (k, c, v1, v2) VALUES (?, ?, ?, ?)"),
                            ConsistencyLevel.ALL, k, k, value, value + 10);

            if (nodesToFlush.length == 0)
                cluster.forEach((instance) -> instance.flush(KEYSPACE));
            else
                cluster.get(nodesToFlush).forEach((instance) -> instance.flush(KEYSPACE));
        }

        return value;
    }

    private void setThresholds(int warn, int fail, int... nodes)
    {
        Stream<IInvokableInstance> instances;

        if (nodes.length == 0)
            instances = cluster.stream();
        else
            instances = cluster.get(nodes).stream();

        instances.forEach(instance -> instance.acceptsOnInstance((IIsolatedExecutor.SerializableBiConsumer<Integer, Integer>)
                                                                 (w, f) -> Guardrails.instance.setSaiSSTableIndexesPerQueryThreshold(w, f))
                                              .accept(warn, fail));
        assertTresholds(warn, fail, nodes);
    }

    private void enableGuardrail()
    {
        cluster.forEach(instance -> instance.runOnInstance((IIsolatedExecutor.SerializableRunnable) () -> Guardrails.instance.setNonPartitionRestrictedQueryEnabled(true)));
        cluster.forEach(instance -> assertTrue(instance.callsOnInstance((IIsolatedExecutor.SerializableCallable<Boolean>) () -> Guardrails.instance.getNonPartitionRestrictedQueryEnabled()).call()));
    }

    private void disableGuardrail()
    {
        cluster.forEach(instance -> instance.runOnInstance((IIsolatedExecutor.SerializableRunnable) () -> Guardrails.instance.setNonPartitionRestrictedQueryEnabled(false)));
        cluster.forEach(instance -> assertFalse(instance.callsOnInstance((IIsolatedExecutor.SerializableCallable<Boolean>) () -> Guardrails.instance.getNonPartitionRestrictedQueryEnabled()).call()));
    }

    private void assertTresholds(int expectedWarn, int expectedFail, int... nodes)
    {
        Stream<IInvokableInstance> instances;

        if (nodes.length == 0)
            instances = cluster.stream();
        else
            instances = cluster.get(nodes).stream();

        instances.forEach(instance -> {
            assertEquals(expectedWarn,
                         instance.callsOnInstance((IIsolatedExecutor.SerializableCallable<Integer>) () -> Guardrails.instance.getSaiSSTableIndexesPerQueryWarnThreshold())
                                 .call().intValue());
            assertEquals(expectedFail,
                         instance.callsOnInstance((IIsolatedExecutor.SerializableCallable<Integer>) () -> Guardrails.instance.getSaiSSTableIndexesPerQueryFailThreshold())
                                 .call().intValue());
        });
    }

    private long totalWarnings()
    {
        return cluster.stream().mapToLong(i -> i.metrics().getCounter("org.apache.cassandra.metrics.keyspace.TooManySSTableIndexesReadWarnings.distributed_test_keyspace")).sum();
    }

    private long totalAborts()
    {
        return cluster.stream().mapToLong(i -> i.metrics().getCounter("org.apache.cassandra.metrics.keyspace.TooManySSTableIndexesReadAborts.distributed_test_keyspace")).sum();
    }
}
