/*
 * 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.bookkeeper.proto.checksum;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;

import com.scurrilous.circe.checksum.IntHash;
import com.scurrilous.circe.checksum.Java8IntHash;
import com.scurrilous.circe.checksum.Java9IntHash;
import com.scurrilous.circe.checksum.JniIntHash;
import com.scurrilous.circe.crc.Sse42Crc32C;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.ByteBufUtil;
import io.netty.buffer.CompositeByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.util.ReferenceCounted;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
import org.apache.bookkeeper.proto.BookieProtoEncoding;
import org.apache.bookkeeper.proto.BookieProtocol;
import org.apache.bookkeeper.util.ByteBufList;
import org.apache.bookkeeper.util.ByteBufVisitor;
import org.apache.commons.lang3.RandomUtils;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;

/**
 * This test class was added to reproduce a bug in the checksum calculation when
 * the payload is a CompositeByteBuf and this buffer has a reader index state other than 0.
 * The reader index state gets lost in the unwrapping process.
 *
 * There were at least 2 different bugs. One that occured when the
 * payload was >= BookieProtoEncoding.SMALL_ENTRY_SIZE_THRESHOLD and the other when
 * it was < BookieProtoEncoding.SMALL_ENTRY_SIZE_THRESHOLD.
 * This test covers both useV2Protocol=true and useV2Protocol=false since the bug was triggered differently.
 *
 * The bug has been fixed and this test is here to make sure it doesn't happen again.
 */
@RunWith(Parameterized.class)
public class CompositeByteBufUnwrapBugReproduceTest {
    final byte[] testPayLoad;
    final int defaultBufferPrefixLength;
    private final boolean useV2Protocol;

    // set to 0 to 3 to run a single scenario for debugging purposes
    private static final int RUN_SINGLE_SCENARIO_FOR_DEBUGGING = -1;

    @Parameterized.Parameters
    public static Collection<Object[]> testScenarios() {
        List<Object[]> scenarios = Arrays.asList(new Object[][] {
                {BookieProtoEncoding.SMALL_ENTRY_SIZE_THRESHOLD - 1, true},
                {BookieProtoEncoding.SMALL_ENTRY_SIZE_THRESHOLD - 1, false},
                {BookieProtoEncoding.SMALL_ENTRY_SIZE_THRESHOLD, true},
                {BookieProtoEncoding.SMALL_ENTRY_SIZE_THRESHOLD, false}
        });
        if (RUN_SINGLE_SCENARIO_FOR_DEBUGGING >= 0) {
            // pick a single scenario for debugging
            scenarios = scenarios.subList(RUN_SINGLE_SCENARIO_FOR_DEBUGGING, 1);
        }
        return scenarios;
    }

    public CompositeByteBufUnwrapBugReproduceTest(int payloadSize, boolean useV2Protocol) {
        this.testPayLoad = createTestPayLoad(payloadSize);
        this.defaultBufferPrefixLength = payloadSize / 7;
        this.useV2Protocol = useV2Protocol;
    }

    private static byte[] createTestPayLoad(int payloadSize) {
        byte[] payload = new byte[payloadSize];
        for (int i = 0; i < payloadSize; i++) {
            payload[i] = (byte) i;
        }
        return payload;
    }


    /**
     * A DigestManager that uses the given IntHash implementation for testing.
     */
    static class TestIntHashDigestManager extends DigestManager {
        private final IntHash intHash;

        public TestIntHashDigestManager(IntHash intHash, long ledgerId, boolean useV2Protocol,
                                        ByteBufAllocator allocator) {
            super(ledgerId, useV2Protocol, allocator);
            this.intHash = intHash;
        }

        @Override
        int getMacCodeLength() {
            return 4;
        }

        @Override
        boolean isInt32Digest() {
            return true;
        }

        @Override
        void populateValueAndReset(int digest, ByteBuf buf) {
            buf.writeInt(digest);
        }

        @Override
        int internalUpdate(int digest, ByteBuf data, int offset, int len) {
            return intHash.resume(digest, data, offset, len);
        }

        @Override
        int internalUpdate(int digest, byte[] buffer, int offset, int len) {
            return intHash.resume(digest, buffer, offset, len);
        }

        @Override
        boolean acceptsMemoryAddressBuffer() {
            return intHash.acceptsMemoryAddressBuffer();
        }
    }

    @Test
    public void shouldCalculateChecksumForCompositeBuffer() {
        ByteBuf testPayload = Unpooled.wrappedBuffer(testPayLoad);
        byte[] referenceOutput = computeDigestAndPackageForSending(new Java8IntHash(), testPayload.retainedDuplicate());
        assertDigestAndPackageMatchesReference(new Java8IntHash(), testPayload, referenceOutput);
        assertDigestAndPackageMatchesReference(new Java9IntHash(), testPayload, referenceOutput);
        if (Sse42Crc32C.isSupported()) {
            assertDigestAndPackageMatchesReference(new JniIntHash(), testPayload, referenceOutput);
        }
        testPayload.release();
    }

    private void assertDigestAndPackageMatchesReference(IntHash intHash, ByteBuf payload, byte[] referenceOutput) {
        assertDigestAndPackageScenario(intHash, payload.retainedDuplicate(), referenceOutput, testPayLoad,
                "plain payload, no wrapping");

        ByteBuf payload2 = wrapWithPrefixAndCompositeByteBufWithReaderIndexState(payload.retainedDuplicate(),
                defaultBufferPrefixLength);
        assertDigestAndPackageScenario(intHash, payload2, referenceOutput, testPayLoad,
                "payload with prefix wrapped in CompositeByteBuf with readerIndex state");

        ByteBuf payload3 = wrapWithPrefixAndMultipleCompositeByteBufWithReaderIndexStateAndMultipleLayersOfDuplicate(
                payload.retainedDuplicate(), defaultBufferPrefixLength);
        assertDigestAndPackageScenario(intHash, payload3, referenceOutput, testPayLoad,
                "payload with prefix wrapped in 2 layers of CompositeByteBuf with readerIndex state in the outer "
                        + "composite. In addition, the outer composite is duplicated twice.");

        ByteBuf payload4 = wrapInCompositeByteBufAndSlice(payload.retainedDuplicate(), defaultBufferPrefixLength);
        assertDigestAndPackageScenario(intHash, payload4, referenceOutput, testPayLoad,
                "payload with prefix wrapped in CompositeByteBuf and sliced");
    }

    private void assertDigestAndPackageScenario(IntHash intHash, ByteBuf payload, byte[] referenceOutput,
                                                byte[] testPayLoadArray,
                                                String scenario) {
        // this validates that the readable bytes in the payload match the TEST_PAYLOAD content
        assertArrayEquals(testPayLoadArray, ByteBufUtil.getBytes(payload.duplicate()),
                "input is invalid for scenario '" + scenario + "'");

        ByteBuf visitedCopy = Unpooled.buffer(payload.readableBytes());
        ByteBufVisitor.visitBuffers(payload, payload.readerIndex(), payload.readableBytes(),
                new ByteBufVisitor.ByteBufVisitorCallback<Void>() {
                    @Override
                    public void visitBuffer(Void context, ByteBuf visitBuffer, int visitIndex, int visitLength) {
                        visitedCopy.writeBytes(visitBuffer, visitIndex, visitLength);
                    }

                    @Override
                    public void visitArray(Void context, byte[] visitArray, int visitIndex, int visitLength) {
                        visitedCopy.writeBytes(visitArray, visitIndex, visitLength);
                    }
                }, null);

        assertArrayEquals(ByteBufUtil.getBytes(visitedCopy), testPayLoadArray,
                "visited copy is invalid for scenario '" + scenario + "'. Bug in ByteBufVisitor?");

        // compute the digest and package
        byte[] output = computeDigestAndPackageForSending(intHash, payload.duplicate());
        if (referenceOutput == null) {
            referenceOutput =
                    computeDigestAndPackageForSending(new Java8IntHash(), Unpooled.wrappedBuffer(testPayLoadArray));
        }
        // this validates that the output matches the reference output
        assertArrayEquals(referenceOutput, output, "output is invalid for scenario '" + scenario + "'");
    }

    private byte[] computeDigestAndPackageForSending(IntHash intHash, ByteBuf data) {
        DigestManager digestManager = new TestIntHashDigestManager(intHash, 1, useV2Protocol, ByteBufAllocator.DEFAULT);
        ReferenceCounted packagedBuffer =
                digestManager.computeDigestAndPackageForSending(1, 0, data.readableBytes(), data,
                        MacDigestManager.EMPTY_LEDGER_KEY, BookieProtocol.FLAG_NONE);
        return packagedBufferToBytes(packagedBuffer);
    }

    ByteBuf wrapWithPrefixAndCompositeByteBufWithReaderIndexState(ByteBuf payload, int bufferPrefixLength) {
        // create a new buffer with a prefix and the actual payload
        ByteBuf prefixedPayload = ByteBufAllocator.DEFAULT.buffer(bufferPrefixLength + payload.readableBytes());
        prefixedPayload.writeBytes(RandomUtils.nextBytes(bufferPrefixLength));
        prefixedPayload.writeBytes(payload);

        // wrap the buffer in a composite buffer
        CompositeByteBuf outerComposite = ByteBufAllocator.DEFAULT.compositeBuffer();
        outerComposite.addComponent(true, prefixedPayload);

        // set reader index state. this is the state that gets lost in the unwrapping process
        outerComposite.readerIndex(bufferPrefixLength);

        return outerComposite;
    }

    ByteBuf wrapWithPrefixAndMultipleCompositeByteBufWithReaderIndexStateAndMultipleLayersOfDuplicate(ByteBuf payload,
                                                                                               int bufferPrefixLength) {
        // create a new buffer with a prefix and the actual payload
        ByteBuf prefixedPayload = ByteBufAllocator.DEFAULT.buffer(bufferPrefixLength + payload.readableBytes());
        prefixedPayload.writeBytes(RandomUtils.nextBytes(bufferPrefixLength));
        prefixedPayload.writeBytes(payload);

        CompositeByteBuf innerComposite = ByteBufAllocator.DEFAULT.compositeBuffer();
        innerComposite.addComponent(true, prefixedPayload);
        innerComposite.addComponent(true, Unpooled.EMPTY_BUFFER);

        // wrap the buffer in a composite buffer
        CompositeByteBuf outerComposite = ByteBufAllocator.DEFAULT.compositeBuffer();
        outerComposite.addComponent(true, innerComposite);
        outerComposite.addComponent(true, Unpooled.EMPTY_BUFFER);

        // set reader index state. this is the state that gets lost in the unwrapping process
        outerComposite.readerIndex(bufferPrefixLength);

        return outerComposite.duplicate().duplicate();
    }

    ByteBuf wrapInCompositeByteBufAndSlice(ByteBuf payload, int bufferPrefixLength) {
        // create a composite buffer
        CompositeByteBuf compositeWithPrefix = ByteBufAllocator.DEFAULT.compositeBuffer();
        compositeWithPrefix.addComponent(true, Unpooled.wrappedBuffer(RandomUtils.nextBytes(bufferPrefixLength)));
        compositeWithPrefix.addComponent(true, payload);

        // return a slice of the composite buffer so that it returns the payload
        return compositeWithPrefix.slice(bufferPrefixLength, payload.readableBytes());
    }

    private static byte[] packagedBufferToBytes(ReferenceCounted packagedBuffer) {
        byte[] output;
        if (packagedBuffer instanceof ByteBufList) {
            ByteBufList bufList = (ByteBufList) packagedBuffer;
            output = new byte[bufList.readableBytes()];
            bufList.getBytes(output);
            for (int i = 0; i < bufList.size(); i++) {
                bufList.getBuffer(i).release();
            }
        } else if (packagedBuffer instanceof ByteBuf) {
            output = ByteBufUtil.getBytes((ByteBuf) packagedBuffer);
            packagedBuffer.release();
        } else {
            throw new RuntimeException("Unexpected type: " + packagedBuffer.getClass());
        }
        return output;
    }
}