/*
 * 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.index.sai.disk.v1.bbtree;

import java.util.function.IntFunction;

import org.junit.Before;
import org.junit.Test;

import org.apache.cassandra.index.sai.disk.v1.segment.SegmentTrieBuffer;
import org.apache.cassandra.index.sai.utils.SAIRandomizedTester;
import org.apache.lucene.store.ByteBuffersDataOutput;
import org.apache.lucene.store.ByteBuffersIndexOutput;
import org.apache.lucene.store.DataInput;

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

public class BlockBalancedTreeTest extends SAIRandomizedTester
{
    private ByteBuffersDataOutput dataOutput;

    @Before
    public void setupDataOutput()
    {
        dataOutput = new ByteBuffersDataOutput();
    }

    @Test
    public void testEmptyTree() throws Exception
    {
        long treeFilePointer = writeBalancedTree(0, 100, rowID -> rowID);

        assertEquals(-1, treeFilePointer);
    }

    @Test
    public void testSingleLeaf() throws Exception
    {
        try (BlockBalancedTreeWalker walker = generateBalancedTree(100, 100, rowID -> rowID))
        {
            assertEquals(1, walker.numLeaves);
            assertEquals(1, walker.treeDepth);
            assertEquals(100, walker.valueCount);

            BlockBalancedTreeWalker.TraversalState state = walker.newTraversalState();

            assertTrue(state.atLeafNode());

            recursiveAssertTraversal(state, -1);

            assertEquals(walker.treeDepth, state.maxLevel + 1);
        }
    }

    @Test
    public void testTreeWithSameValue() throws Exception
    {
        try (BlockBalancedTreeWalker walker = generateBalancedTree(100, 4, rowID -> 1))
        {
            BlockBalancedTreeWalker.TraversalState state = walker.newTraversalState();

            recursiveAssertTraversal(state, -1);

            assertEquals(walker.treeDepth, state.maxLevel + 1);
        }
    }

    @Test
    public void testTreeDepthNeverMoreThanNumberOfLeaves() throws Exception
    {
        int leafSize = 4;
        for (int numLeaves = 1; numLeaves < 1000; numLeaves++)
        {
            int numRows = leafSize * numLeaves;

            try (BlockBalancedTreeWalker walker = generateBalancedTree(numRows, leafSize, rowID -> rowID))
            {
                assertEquals(numLeaves, walker.numLeaves);
                assertTrue(walker.treeDepth <= walker.numLeaves);

                BlockBalancedTreeWalker.TraversalState state = walker.newTraversalState();

                recursiveAssertTraversal(state, -1);

                assertEquals(walker.treeDepth, state.maxLevel + 1);
            }
        }
    }

    @Test
    public void randomisedTreeTest() throws Exception
    {
        int loops = nextInt(10, 1000);

        for (int loop = 0; loop < loops; loop++)
        {
            int leafSize = nextInt(2, 512);
            int numRows = nextInt(1000, 10000);

            try (BlockBalancedTreeWalker walker = generateBalancedTree(numRows, leafSize, rowID -> nextInt(0, numRows / 2)))
            {
                BlockBalancedTreeWalker.TraversalState state = walker.newTraversalState();

                recursiveAssertTraversal(state, -1);

                assertEquals(walker.treeDepth, state.maxLevel + 1);
            }
        }
    }

    private long recursiveAssertTraversal(BlockBalancedTreeWalker.TraversalState state, long lastLeafBlockFP)
    {
        if (state.atLeafNode())
        {
            assertTrue(state.nodeExists());
            assertTrue(state.getLeafBlockFP() > lastLeafBlockFP);
            return state.getLeafBlockFP();
        }
        else
        {
            state.pushLeft();
            lastLeafBlockFP = recursiveAssertTraversal(state, lastLeafBlockFP);
            state.pop();

            state.pushRight();
            lastLeafBlockFP = recursiveAssertTraversal(state, lastLeafBlockFP);
            state.pop();

            return lastLeafBlockFP;
        }
    }

    private BlockBalancedTreeWalker generateBalancedTree(int numRows, int leafSize, IntFunction<Integer> valueProvider) throws Exception
    {
        long treeOffset = writeBalancedTree(numRows, leafSize, valueProvider);

        DataInput input = dataOutput.toDataInput();

        return new BlockBalancedTreeWalker(input, treeOffset);
    }

    private long writeBalancedTree(int numRows, int leafSize, IntFunction<Integer> valueProvider) throws Exception
    {
        SegmentTrieBuffer buffer = new SegmentTrieBuffer();

        for (int rowId = 0; rowId < numRows; rowId++)
        {
            buffer.add(integerToByteComparable(valueProvider.apply(rowId)), Integer.BYTES, rowId);
        }

        BlockBalancedTreeWriter writer = new BlockBalancedTreeWriter(4, leafSize);
        ByteBuffersIndexOutput output = new ByteBuffersIndexOutput(dataOutput, "test", "test");
        return writer.write(output, buffer.iterator(), (leafPostings, offset, count) -> {});
    }
}
