/*
 * 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.db.compaction.unified;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.stream.Collectors;

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

import org.apache.cassandra.cql3.CQLTester;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.compaction.CompactionInfo;
import org.apache.cassandra.db.compaction.CompactionManager;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.utils.TimeUUID;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

@RunWith(BMUnitRunner.class)
@BMUnitConfig(debug=true)
@BMRule(
name = "Get stats before task completion",
targetClass = "org.apache.cassandra.db.compaction.ActiveCompactions",
targetMethod = "finishCompaction",
targetLocation = "AT ENTRY",
action = "org.apache.cassandra.db.compaction.unified.BackgroundCompactionTrackingTest.getStats()"
)
public class BackgroundCompactionTrackingTest extends CQLTester
{
    // Get rid of commitlog noise
    @Before
    public void disableCommitlog()
    {
        schemaChange("ALTER KEYSPACE " + KEYSPACE + " WITH durable_writes = false");
    }
    @After
    public void enableCommitlog()
    {
        schemaChange("ALTER KEYSPACE " + KEYSPACE + " WITH durable_writes = true");
    }

    @Test
    public void testBackgroundCompactionTracking()
    {
        testBackgroundCompactionTracking(false, 3);
    }

    @Test
    public void testBackgroundCompactionTrackingParallelized()
    {
        testBackgroundCompactionTracking(true, 3);
    }

    public void testBackgroundCompactionTracking(boolean parallelize, int shards)
    {
        int bytes_per_row = 5000;
        int partitions = 5000;
        int rows_per_partition = 10;
        long targetSize = 1L * partitions * rows_per_partition * bytes_per_row / shards;
        CompactionManager.instance.setMaximumCompactorThreads(50);
        CompactionManager.instance.setCoreCompactorThreads(50);
        String table = createTable(String.format("CREATE TABLE %%s (k int, t int, v blob, PRIMARY KEY (k, t))" +
                                                 " with compression = {'enabled': false} " +
                                                 " and compaction = {" +
                                                 "'class': 'UnifiedCompactionStrategy', " +
                                                 "'parallelize_output_shards': '%s', " +
                                                 "'base_shard_count': %d, " +
                                                 "'min_sstable_size': '%dB', " + // to disable sharding on flush
                                                 "'scaling_parameters': 'T4, T7'" +
                                   "}",
                                   parallelize, shards, targetSize));
        ColumnFamilyStore cfs = getColumnFamilyStore(KEYSPACE, table);
        cfs.disableAutoCompaction();

        for (int iter = 1; iter <= 5; ++iter)
        {
            byte [] payload = new byte[bytes_per_row];
            new Random(42).nextBytes(payload);
            ByteBuffer b = ByteBuffer.wrap(payload);
            Set<SSTableReader> before = new HashSet<>(cfs.getLiveSSTables());

            for (int i = 0; i < partitions; i++)
            {
                for (int j = 0; j < rows_per_partition; j++)
                    execute(String.format("INSERT INTO %s.%s(k, t, v) VALUES (?, ?, ?)", KEYSPACE, table), i, j, b);

                if ((i + 1) % ((partitions  + 3) / 4) == 0)
                    cfs.forceBlockingFlush(ColumnFamilyStore.FlushReason.UNIT_TESTS);
            }
            cfs.forceBlockingFlush(ColumnFamilyStore.FlushReason.UNIT_TESTS);
            operations = new ArrayList<>();
            Set<SSTableReader> newSSTables = new HashSet<>(cfs.getLiveSSTables());
            newSSTables.removeAll(before);
            long uncompressedSize = newSSTables.stream().mapToLong(SSTableReader::uncompressedLength).sum();

            cfs.enableAutoCompaction(true); // since the trigger is hit, this initiates an L0 compaction
            cfs.disableAutoCompaction();

            // Check that the background compactions state is correct during the compaction
            Assert.assertTrue("Byteman rule did not fire", !operations.isEmpty());
            printStats();
            int tasks = parallelize ? shards : 1;
            assertEquals(tasks, operations.size());
            TimeUUID mainOpId = null;
            for (int i = 0; i < operations.size(); ++i)
            {
                BitSet seqs = new BitSet(shards);
                int expectedSize = tasks - i;
                final List<CompactionInfo> ops = getCompactionOps(i, cfs);
                final int size = ops.size();
                int finished = tasks - size;
                assertTrue(size >= expectedSize); // some task may have not managed to close
                for (var op : ops)
                {
                    final TimeUUID opIdSeq0 = op.getTaskId().withSequence(0);
                    if (mainOpId == null)
                        mainOpId = opIdSeq0;
                    else
                        assertEquals(mainOpId, opIdSeq0);
                    seqs.set(op.getTaskId().sequence());

                    if (i == 0)
                        Assert.assertEquals(uncompressedSize * 1.0 / tasks, op.getTotal(), uncompressedSize * 0.03);
                    assertTrue("Expected:" + op.getCompleted() +"<=" + (long)(op.getTotal() * 1.03), op.getCompleted() <= op.getTotal() * 1.03);
                    if (op.getCompleted() >= op.getTotal() * 0.97)
                        ++finished;
                }
                assertTrue(finished > i);
                assertEquals(size, seqs.cardinality());
            }
            assertEquals(iter * shards, cfs.getLiveSSTables().size());

            // Check that the background compactions state is correct after the compaction
            operations.clear();
            getStats();
            printStats();
            final List<CompactionInfo> ops = getCompactionOps(0, cfs);
            assertEquals(0, ops.size());
        }
    }

    private static List<CompactionInfo> getCompactionOps(int i, ColumnFamilyStore cfs)
    {
        return operations.get(i)
                         .stream()
                         .filter(op -> op.getTableMetadata() == cfs.metadata())
                         .collect(Collectors.toList());
    }

    private void printStats()
    {
        for (int i = 0; i < operations.size(); ++i)
        {
            System.out.println(operations.get(i).stream().map(Object::toString).collect(Collectors.joining("\n")));
        }
    }

    public static synchronized void getStats()
    {
        operations.add(CompactionManager.instance.getSSTableTasks());
    }

    static List<List<CompactionInfo>> operations;
}
