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

import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.jboss.byteman.contrib.bmunit.BMRule;
import org.jboss.byteman.contrib.bmunit.BMRules;
import org.jboss.byteman.contrib.bmunit.BMUnitRunner;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;

import org.apache.cassandra.db.ReadQuery;
import org.apache.cassandra.db.SystemKeyspace;
import org.apache.cassandra.db.marshal.Int32Type;
import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.schema.SchemaConstants;
import org.apache.cassandra.schema.SchemaKeyspaceTables;
import org.apache.cassandra.schema.TableMetadata;
import org.apache.cassandra.service.ClientState;
import org.apache.cassandra.transport.Dispatcher;
import org.apache.cassandra.utils.ByteBufferUtil;
import org.apache.cassandra.utils.MD5Digest;

import static java.util.Collections.emptyMap;
import static org.apache.cassandra.service.QueryState.forInternalCalls;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

@RunWith(BMUnitRunner.class)
public class PstmtPersistenceTest extends CQLTester
{
    private static final CompletableFuture<?>[] futureArray = new CompletableFuture[0];

    private static final ConcurrentMap<MD5Digest, Long> preparedStatementLoadTimestamps = new ConcurrentHashMap<>();
    private static final ConcurrentMap<MD5Digest, Long> preparedStatementRemoveTimestamps = new ConcurrentHashMap<>();

    // page size passed to preloadPreparedStatements
    private static final int PRELOAD_PAGE_SIZE = 100;

    // recorded page invocations in preloadPreparedStatements
    private static final AtomicInteger pageInvocations = new AtomicInteger();

    @Before
    public void setUp()
    {
        preparedStatementLoadTimestamps.clear();
        preparedStatementRemoveTimestamps.clear();

        QueryProcessor.clearPreparedStatements(false);
    }

    @Test
    public void testCachedPreparedStatements() throws Throwable
    {
        // need this for pstmt execution/validation tests
        requireNetwork();

        assertEquals(0, numberOfStatementsOnDisk());

        execute("CREATE KEYSPACE IF NOT EXISTS foo WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'}");
        execute("CREATE TABLE foo.bar (key text PRIMARY KEY, val int)");

        ClientState clientState = ClientState.forExternalCalls(InetSocketAddress.createUnresolved("127.0.0.1", 1234));

        createTable("CREATE TABLE %s (pk int PRIMARY KEY, val text)");

        List<MD5Digest> stmtIds = new ArrayList<>();
        String statement0 = "SELECT * FROM %s WHERE keyspace_name = ?";
        String statement1 = "SELECT * FROM %s WHERE pk = ?";
        String statement2 = "SELECT * FROM %s WHERE key = ?";
        String statement3 = "SELECT * FROM %S WHERE key = ?";
        stmtIds.add(prepareStatement(statement0, SchemaConstants.SCHEMA_KEYSPACE_NAME, SchemaKeyspaceTables.TABLES, clientState));
        stmtIds.add(prepareStatement(statement1, clientState));
        stmtIds.add(prepareStatement(statement2, "foo", "bar", clientState));
        clientState.setKeyspace("foo");

        stmtIds.add(prepareStatement(statement1, clientState));
        stmtIds.add(prepareStatement(statement3, "foo", "bar", clientState));

        assertEquals(5, stmtIds.size());
        // statement1 will have two statements prepared because of `setKeyspace` usage
        assertEquals(6, QueryProcessor.preparedStatementsCount());
        assertEquals(6, numberOfStatementsOnDisk());

        QueryHandler handler = ClientState.getCQLQueryHandler();
        validatePstmts(stmtIds, handler);

        // clear prepared statements cache
        QueryProcessor.clearPreparedStatements(true);
        assertEquals(0, QueryProcessor.preparedStatementsCount());
        for (MD5Digest stmtId : stmtIds)
            Assert.assertNull(handler.getPrepared(stmtId));

        // load prepared statements and validate that these still execute fine
        QueryProcessor.instance.preloadPreparedStatements();
        validatePstmts(stmtIds, handler);


        // validate that the prepared statements are in the system table
        String queryAll = "SELECT * FROM " + SchemaConstants.SYSTEM_KEYSPACE_NAME + '.' + SystemKeyspace.PREPARED_STATEMENTS;
        for (UntypedResultSet.Row row : QueryProcessor.executeOnceInternal(queryAll))
        {
            MD5Digest digest = MD5Digest.wrap(ByteBufferUtil.getArray(row.getBytes("prepared_id")));
            QueryProcessor.Prepared prepared = QueryProcessor.instance.getPrepared(digest);
            Assert.assertNotNull(prepared);
        }

        // add another prepared statement and sync it to table
        prepareStatement(statement2, "foo", "bar", clientState);

        // statement1 will have two statements prepared because of `setKeyspace` usage
        assertEquals(7, numberOfStatementsInMemory());
        assertEquals(7, numberOfStatementsOnDisk());

        // drop a keyspace (prepared statements are removed - syncPreparedStatements() remove should the rows, too)
        execute("DROP KEYSPACE foo");
        assertEquals(3, numberOfStatementsInMemory());
        assertEquals(3, numberOfStatementsOnDisk());
    }

    private void validatePstmts(List<MD5Digest> stmtIds, QueryHandler handler)
    {
        QueryOptions optionsStr = QueryOptions.forInternalCalls(Collections.singletonList(UTF8Type.instance.fromString("foobar")));
        QueryOptions optionsInt = QueryOptions.forInternalCalls(Collections.singletonList(Int32Type.instance.decompose(42)));
        validatePstmt(handler, stmtIds.get(0), optionsStr);
        validatePstmt(handler, stmtIds.get(1), optionsInt);
        validatePstmt(handler, stmtIds.get(2), optionsStr);
        validatePstmt(handler, stmtIds.get(3), optionsInt);
        validatePstmt(handler, stmtIds.get(4), optionsStr);
    }

    private static void validatePstmt(QueryHandler handler, MD5Digest stmtId, QueryOptions options)
    {
        QueryProcessor.Prepared prepared = handler.getPrepared(stmtId);
        Assert.assertNotNull(prepared);
        handler.processPrepared(prepared.statement, forInternalCalls(), options, emptyMap(), Dispatcher.RequestTime.forImmediateExecution());
    }

    @Test
    public void testPstmtInvalidation() throws Throwable
    {
        ClientState clientState = ClientState.forInternalCalls();

        createTable("CREATE TABLE %s (key int primary key, val int)");

        long initialEvicted = numberOfEvictedStatements();

        for (int cnt = 1; cnt < 10000; cnt++)
        {
            prepareStatement("INSERT INTO %s (key, val) VALUES (?, ?) USING TIMESTAMP " + cnt, clientState);

            if (numberOfEvictedStatements() - initialEvicted > 0)
            {
                assertEquals("Number of statements in table and in cache don't match", numberOfStatementsInMemory(), numberOfStatementsOnDisk());

                // prepare more statements to trigger more evictions
                for (int cnt2 = cnt + 1; cnt2 < cnt + 10; cnt2++)
                    prepareStatement("INSERT INTO %s (key, val) VALUES (?, ?) USING TIMESTAMP " + cnt2, clientState);

                // each new prepared statement should have caused an eviction
                assertEquals("eviction count didn't increase by the expected number", 10, numberOfEvictedStatements() - initialEvicted);
                assertEquals("Number of statements in memory (expected) and table (actual) don't match", numberOfStatementsInMemory(), numberOfStatementsOnDisk());

                return;
            }
        }

        fail("Prepared statement eviction does not work");
    }

    @Test
    @BMRules(rules= {
             @BMRule(name = "CaptureWriteTimestamps",
                     targetClass = "SystemKeyspace",
                     targetMethod = "writePreparedStatement(String, MD5Digest, String, long)",
                     targetLocation = "AT INVOKE executeInternal",
                     action = "org.apache.cassandra.cql3.PstmtPersistenceTest.preparedStatementLoadTimestamps.put($key, $timestamp);"
             ),
             @BMRule(name = "CaptureEvictTimestamps",
                     targetClass = "QueryProcessor",
                     targetMethod = "evictPreparedStatement(MD5Digest, RemovalCause)",
                     action = "org.apache.cassandra.cql3.PstmtPersistenceTest.preparedStatementRemoveTimestamps.put($key, org.apache.cassandra.service.ClientState.getTimestamp());"
             )
    })
    public void testAsyncPstmtInvalidation() throws Throwable
    {
        ClientState clientState = ClientState.forInternalCalls();
        createTable("CREATE TABLE %s (key int primary key, val int)");

        // prepare statements concurrently in a thread pool to exercise bug encountered in CASSANDRA-19703 where
        // delete from table occurs before the insert due to early eviction.
        final ExecutorService executor = Executors.newFixedThreadPool(10);

        long initialEvicted = numberOfEvictedStatements();
        try
        {
            int initialMaxStatementsToPrepare = 10000;
            int maxStatementsToPrepare = initialMaxStatementsToPrepare;
            boolean hasEvicted = false;
            int concurrency = 100;
            List<CompletableFuture<MD5Digest>> prepareFutures = new ArrayList<>(concurrency);

            for (int cnt = 1; cnt <= maxStatementsToPrepare; cnt++)
            {
                final int localCnt = cnt;
                prepareFutures.add(CompletableFuture.supplyAsync(() -> prepareStatement("INSERT INTO %s (key, val) VALUES (?, ?) USING TIMESTAMP " + localCnt, clientState), executor));

                if (prepareFutures.size() == concurrency)
                {
                    // Await completion of current inflight futures
                    CompletableFuture.allOf(prepareFutures.toArray(futureArray)).get(10, TimeUnit.SECONDS);
                    prepareFutures.clear();
                }

                // Once we've detected evictions, prepare as many statements as we've prepared so far to initialMaxStatementsToPrepare and then stop.
                if (!hasEvicted && numberOfEvictedStatements() - initialEvicted > 0)
                {
                    maxStatementsToPrepare = Math.min(cnt * 2, initialMaxStatementsToPrepare);
                    hasEvicted = true;
                }
            }

            long evictedStatements = numberOfEvictedStatements() - initialEvicted;
            assertNotEquals("Should have evicted some prepared statements", 0, evictedStatements);

            // Recorded prepared statement removals should match metrics
            assertEquals("Actual evicted statements does not match metrics", evictedStatements, preparedStatementRemoveTimestamps.size());

            // For each prepared statement evicted, assert the time it was deleted is greater than the timestamp
            // used for when it was loaded.
            for (Map.Entry<MD5Digest, Long> evictedStatementEntry : preparedStatementRemoveTimestamps.entrySet())
            {
                MD5Digest key = evictedStatementEntry.getKey();
                long deletionTimestamp = evictedStatementEntry.getValue();
                long insertionTimestamp = preparedStatementLoadTimestamps.get(key);

                assertTrue(String.format("Expected deletion timestamp for prepared statement (%d) to be greater than insertion timestamp (%d)",
                                         deletionTimestamp, insertionTimestamp),
                           deletionTimestamp > insertionTimestamp);
            }

            // ensure the number of statements on disk match the number in memory, if number of statements on disk eclipses in memory, there was a leak.
            assertEquals("Number of statements in memory (expected) and table (actual) don't match", numberOfStatementsInMemory(), numberOfStatementsOnDisk());
        }
        finally
        {
            executor.shutdown();
        }
    }

    /**
     * Invoked whenever paging happens in testPreloadPreparedStatements, increments PAGE_INVOCATIONS when we detect
     * paging happening in the path of QueryProcessor.preloadPreparedStatements with the expected page size.
     */
    @SuppressWarnings("unused")
    private static void nextPageReadQuery(ReadQuery query, int pageSize)
    {
        TableMetadata metadata = query.metadata();
        if (metadata.keyspace.equals(SchemaConstants.SYSTEM_KEYSPACE_NAME) &&
            metadata.name.equals(SystemKeyspace.PREPARED_STATEMENTS) &&
            pageSize == PRELOAD_PAGE_SIZE)
        {
            for (StackTraceElement stackTraceElement : Thread.currentThread().getStackTrace())
            {
                if (stackTraceElement.getClassName().equals(QueryProcessor.class.getName()) && stackTraceElement.getMethodName().equals("preloadPreparedStatements"))
                {
                    pageInvocations.incrementAndGet();
                    return;
                }
            }
        }
    }

    @Test
    @BMRule(name = "CapturePageInvocations",
            targetClass = "PartitionRangeQueryPager",
            targetMethod = "nextPageReadQuery(int)",
            action = "org.apache.cassandra.cql3.PstmtPersistenceTest.nextPageReadQuery($this.query, $pageSize)")
    public void testPreloadPreparedStatements() throws Throwable
    {
        ClientState clientState = ClientState.forInternalCalls();
        createTable("CREATE TABLE %s (key int primary key, val int)");

        // Prepare more statements than the paging size to ensure paging works properly.
        int statementsToPrepare = 750;

        for (int cnt = 1; cnt <= statementsToPrepare; cnt++)
        {
            prepareStatement("INSERT INTO %s (key, val) VALUES (?, ?) USING TIMESTAMP " + cnt, clientState);
        }

        // Capture how many statements are in memory before clearing cache.
        long statementsInMemory = numberOfStatementsInMemory();
        long statementsOnDisk = numberOfStatementsOnDisk();
        assertEquals(statementsOnDisk, statementsInMemory);

        // Drop prepared statements from cache only and ensure the cache empties out.
        QueryProcessor.clearPreparedStatements(true);
        assertEquals(0, numberOfStatementsInMemory());

        // Load prepared statements and ensure the cache size matches max
        QueryProcessor.instance.preloadPreparedStatements(PRELOAD_PAGE_SIZE);

        long statementsInMemoryAfterLoading = numberOfStatementsInMemory();
        // Ensure size of cache matches statements that were on disk before preload
        assertEquals("Statements prepared - evicted (expected) does not match statements in memory (actual)",
                     statementsOnDisk, statementsInMemoryAfterLoading);

        // Number of statements on disk shold match memory
        assertEquals(statementsInMemoryAfterLoading, numberOfStatementsOnDisk());

        // Ensure only executed the expected amount of pages.
        int expectedPageInvocations = (int) Math.ceil(statementsInMemoryAfterLoading / (double) PRELOAD_PAGE_SIZE);
        assertEquals(expectedPageInvocations, pageInvocations.get());
    }

    @Test
    public void testPreloadPreparedStatementsUntilCacheFull()
    {
        QueryHandler handler = ClientState.getCQLQueryHandler();
        ClientState clientState = ClientState.forInternalCalls();
        createTable("CREATE TABLE %s (key int primary key, val int)");

        // Fill up and clear the prepared statement cache several times to load up the system.prepared_statements table.
        // This simulates a 'leak' of prepared statements akin to CASSANDRA-19703 as the system.prepared_statements
        // table is able to grow to a larger size than the in memory prepared statement cache.  In such a case we
        // should detect a possible leak and defer paging indefinitely by returning early in preloadPreparedStatements.
        int statementsLoadedWhenFull = -1;
        long accumulatedSize = 0;
        // load enough prepared statements to fill the cache 5 times.
        for (int cnt = 0; accumulatedSize < QueryProcessor.PREPARED_STATEMENT_CACHE_SIZE_BYTES * 5; cnt++)
        {
            MD5Digest id = prepareStatement("INSERT INTO %s (key, val) VALUES (?, ?) USING TIMESTAMP " + cnt, clientState);
            QueryHandler.Prepared prepared = handler.getPrepared(id);
            assertTrue(prepared.pstmntSize > -1);
            accumulatedSize += prepared.pstmntSize;
            if (statementsLoadedWhenFull == -1 && accumulatedSize > QueryProcessor.PREPARED_STATEMENT_CACHE_SIZE_BYTES)
            {
                statementsLoadedWhenFull = cnt;
            }
            // clear cache repeatedly to avoid eviction.
            QueryProcessor.clearPreparedStatements(true);
        }


        int preloadedStatements = QueryProcessor.instance.preloadPreparedStatements(PRELOAD_PAGE_SIZE);

        // Should have loaded as many statements as we detected were loaded before cache would be full.
        assertTrue(String.format("Preloaded %d statements, expected at least %d",
                                 preloadedStatements, statementsLoadedWhenFull),
                   preloadedStatements > statementsLoadedWhenFull);

        // We should only expect to load how many statements we were able to load before filling the cache
        // + a buffer of 110%, set to 1.5x just to deal with sensitivity of detecting cache filling up.
        int atMostPreloadedExpected = (int) (statementsLoadedWhenFull * 1.5);
        assertTrue(String.format("Preloaded %d statements, but only expected that we'd load at most %d",
                                 preloadedStatements, atMostPreloadedExpected),
                   preloadedStatements <= atMostPreloadedExpected);
    }

    private long numberOfStatementsOnDisk() throws Throwable
    {
        UntypedResultSet.Row row = execute("SELECT COUNT(*) FROM " + SchemaConstants.SYSTEM_KEYSPACE_NAME + '.' + SystemKeyspace.PREPARED_STATEMENTS).one();
        return row.getLong("count");
    }

    private long numberOfStatementsInMemory()
    {
        return QueryProcessor.preparedStatementsCount();
    }

    private long numberOfEvictedStatements()
    {
        return QueryProcessor.metrics.preparedStatementsEvicted.getCount();
    }

    private MD5Digest prepareStatement(String stmt, ClientState clientState)
    {
        return prepareStatement(stmt, keyspace(), currentTable(), clientState);
    }

    private MD5Digest prepareStatement(String stmt, String keyspace, String table, ClientState clientState)
    {
        return QueryProcessor.instance.prepare(String.format(stmt, keyspace + '.' + table), clientState).statementId;
    }
}
