/*
 * 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.test.microbench.sai;

import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.util.Arrays;
import java.util.concurrent.TimeUnit;

import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Threads;
import org.openjdk.jmh.annotations.Warmup;

import org.apache.cassandra.Util;
import org.apache.cassandra.config.CassandraRelevantProperties;
import org.apache.cassandra.config.DatabaseDescriptor;
import org.apache.cassandra.db.Clustering;
import org.apache.cassandra.db.DecoratedKey;
import org.apache.cassandra.db.marshal.CompositeType;
import org.apache.cassandra.db.marshal.LongType;
import org.apache.cassandra.db.marshal.UTF8Type;
import org.apache.cassandra.dht.Murmur3Partitioner;
import org.apache.cassandra.index.sai.disk.PrimaryKeyMap;
import org.apache.cassandra.index.sai.disk.format.IndexDescriptor;
import org.apache.cassandra.index.sai.disk.v1.SSTableComponentsWriter;
import org.apache.cassandra.index.sai.disk.v1.WidePrimaryKeyMap;
import org.apache.cassandra.index.sai.utils.PrimaryKey;
import org.apache.cassandra.io.sstable.Descriptor;
import org.apache.cassandra.io.sstable.format.SSTableReader;
import org.apache.cassandra.io.util.File;
import org.apache.cassandra.schema.TableMetadata;

import static org.apache.cassandra.index.sai.SAITester.getRandom;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

@BenchmarkMode({Mode.Throughput})
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@Warmup(iterations = 3, time = 1)
@Measurement(iterations = 3, time = 5)
@Fork(value = 1, jvmArgsAppend = "-Xmx512M")
@Threads(1)
@State(Scope.Benchmark)
public class KeyLookupBench
{
    private static final int rows = 1_000_000;

    static
    {
        DatabaseDescriptor.toolInitialization();
        // Partitioner is not set in client mode.
        if (DatabaseDescriptor.getPartitioner() == null)
            DatabaseDescriptor.setPartitionerUnsafe(Murmur3Partitioner.instance);
    }

    protected TableMetadata metadata;
    protected IndexDescriptor indexDescriptor;

    private PrimaryKeyMap primaryKeyMap;

    private PrimaryKey primaryKey;

    @Param({"3", "4", "5"})
    public int partitionBlockShift;

    @Param({"3", "4", "5"})
    public int clusteringBlockShift;

    @Param({"10", "100", "1000", "10000"})
    public int partitionSize;

    @Param({"true", "false"})
    public boolean randomClustering;

    @Setup(Level.Trial)
    public void trialSetup() throws Exception
    {
        String keyspaceName = "ks";
        String tableName = this.getClass().getSimpleName();
        metadata = TableMetadata
                   .builder(keyspaceName, tableName)
                   .partitioner(Murmur3Partitioner.instance)
                   .addPartitionKeyColumn("pk1", LongType.instance)
                   .addPartitionKeyColumn("pk2", LongType.instance)
                   .addClusteringColumn("ck1", UTF8Type.instance)
                   .addClusteringColumn("ck2", UTF8Type.instance)
                   .build();

        Descriptor descriptor = new Descriptor(new File(Files.createTempDirectory("jmh").toFile()),
                                               metadata.keyspace,
                                               metadata.name,
                                               Util.newUUIDGen().get());

        indexDescriptor = IndexDescriptor.create(descriptor, metadata.partitioner, metadata.comparator);

        CassandraRelevantProperties.SAI_SORTED_TERMS_PARTITION_BLOCK_SHIFT.setInt(partitionBlockShift);
        CassandraRelevantProperties.SAI_SORTED_TERMS_CLUSTERING_BLOCK_SHIFT.setInt(clusteringBlockShift);
        SSTableComponentsWriter writer = new SSTableComponentsWriter(indexDescriptor);

        PrimaryKey.Factory factory = new PrimaryKey.Factory(metadata.partitioner, metadata.comparator);

        PrimaryKey[] primaryKeys = new PrimaryKey[rows];
        int partition = 0;
        int partitionRowCounter = 0;
        for (int index = 0; index < rows; index++)
        {
            primaryKeys[index] = factory.create(makeKey(metadata, (long) partition, (long) partition), makeClustering(metadata));
            partitionRowCounter++;
            if (partitionRowCounter == partitionSize)
            {
                partition++;
                partitionRowCounter = 0;
            }
        }

        Arrays.sort(primaryKeys);

        DecoratedKey lastKey = null;
        for (PrimaryKey primaryKey : primaryKeys)
        {
            if (lastKey == null || lastKey.compareTo(primaryKey.partitionKey()) < 0)
            {
                lastKey = primaryKey.partitionKey();
                writer.startPartition(lastKey);
            }
            writer.nextRow(primaryKey);
        }

        writer.complete();

        SSTableReader sstableReader = mock(SSTableReader.class);
        when(sstableReader.metadata()).thenReturn(metadata);

        PrimaryKeyMap.Factory mapFactory = new WidePrimaryKeyMap.Factory(indexDescriptor, sstableReader);

        primaryKeyMap = mapFactory.newPerSSTablePrimaryKeyMap();

        primaryKey = primaryKeys[500000];
    }

    @Benchmark
    public long advanceToKey()
    {
        return primaryKeyMap.rowIdFromPrimaryKey(primaryKey);
    }

    private static DecoratedKey makeKey(TableMetadata table, Object...partitionKeys)
    {
        ByteBuffer key;
        if (table.partitionKeyType instanceof CompositeType)
            key = ((CompositeType)table.partitionKeyType).decompose(partitionKeys);
        else
            key = table.partitionKeyType.fromString((String)partitionKeys[0]);
        return table.partitioner.decorateKey(key);
    }

    private Clustering<?> makeClustering(TableMetadata table)
    {
        Clustering<?> clustering;
        if (table.comparator.size() == 0)
            clustering = Clustering.EMPTY;
        else
        {
            ByteBuffer[] values = new ByteBuffer[table.comparator.size()];
            for (int index = 0; index < table.comparator.size(); index++)
                values[index] = table.comparator.subtype(index).fromString(makeClusteringString());
            clustering = Clustering.make(values);
        }
        return clustering;
    }

    private String makeClusteringString()
    {
        if (randomClustering)
            return getRandom().nextTextString(10, 100);
        else
            return String.format("%08d", getRandom().nextIntBetween(0, partitionSize));
    }
}
