/*
 * 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.kafka.common.utils;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import org.junit.jupiter.params.provider.ValueSource;

import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;
import java.util.stream.Stream;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.spy;

public class ChunkedBytesStreamTest {
    private static final Random RANDOM = new Random(1337);
    private final BufferSupplier supplier = BufferSupplier.NO_CACHING;

    @Test
    public void testEofErrorForMethodReadFully() throws IOException {
        ByteBuffer input = ByteBuffer.allocate(8);
        int lengthGreaterThanInput = input.capacity() + 1;
        byte[] got = new byte[lengthGreaterThanInput];
        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(input), supplier, 10, false)) {
            assertEquals(8, is.read(got, 0, got.length), "Should return 8 signifying end of input");
        }
    }

    @ParameterizedTest
    @MethodSource("provideSourceBytebuffersForTest")
    public void testCorrectnessForMethodReadFully(ByteBuffer input) throws IOException {
        byte[] got = new byte[input.array().length];
        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(input), supplier, 10, false)) {
            // perform a 2 pass read. this tests the scenarios where one pass may lead to partially consumed
            // intermediate buffer
            int toRead = RANDOM.nextInt(got.length);
            is.read(got, 0, toRead);
            is.read(got, toRead, got.length - toRead);
        }
        assertArrayEquals(input.array(), got);
    }

    @ParameterizedTest
    @MethodSource("provideCasesWithInvalidInputsForMethodRead")
    public void testInvalidInputsForMethodRead(byte[] b, int off, int len) throws IOException {
        ByteBuffer buffer = ByteBuffer.allocate(16);

        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(buffer.duplicate()), supplier, 10, false)) {
            assertThrows(IndexOutOfBoundsException.class, () -> is.read(b, off, len));
        }
    }

    @ParameterizedTest
    @MethodSource("provideSourceBytebuffersForTest")
    public void testCorrectnessForMethodReadByte(ByteBuffer input) throws IOException {
        byte[] got = new byte[input.array().length];
        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(input), supplier, 10, false)) {
            int i = 0;
            while (i < got.length) {
                got[i++] = (byte) is.read();
            }
        }
        assertArrayEquals(input.array(), got);
    }

    @ParameterizedTest
    @MethodSource("provideSourceBytebuffersForTest")
    public void testCorrectnessForMethodRead(ByteBuffer inputBuf) throws IOException {
        int[] inputArr = new int[inputBuf.capacity()];
        for (int i = 0; i < inputArr.length; i++) {
            inputArr[i] = Byte.toUnsignedInt(inputBuf.get());
        }
        int[] got = new int[inputArr.length];
        inputBuf.rewind();
        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(inputBuf), supplier, 10, false)) {
            int i = 0;
            while (i < got.length) {
                got[i++] = is.read();
            }
        }
        assertArrayEquals(inputArr, got);
    }

    @Test
    public void testEndOfFileForMethodRead() throws IOException {
        ByteBuffer inputBuf = ByteBuffer.allocate(2);
        int lengthGreaterThanInput = inputBuf.capacity() + 1;

        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(inputBuf), supplier, 10, false)) {
            int cnt = 0;
            while (cnt++ < lengthGreaterThanInput) {
                int res = is.read();
                if (cnt > inputBuf.capacity())
                    assertEquals(-1, res, "end of file for read should be -1");
            }
        }
    }

    @ParameterizedTest
    @ValueSource(booleans = {true, false})
    public void testEndOfSourceForMethodSkip(boolean pushSkipToSourceStream) throws IOException {
        ByteBuffer inputBuf = ByteBuffer.allocate(16);
        RANDOM.nextBytes(inputBuf.array());
        inputBuf.rewind();

        final InputStream sourcestream = spy(new ByteBufferInputStream(inputBuf));
        try (InputStream is = new ChunkedBytesStream(sourcestream, supplier, 10, pushSkipToSourceStream)) {
            long res = is.skip(inputBuf.capacity() + 1);
            assertEquals(inputBuf.capacity(), res);
        }
    }

    @ParameterizedTest
    @MethodSource({"provideSourceSkipValuesForTest"})
    public void testCorrectnessForMethodSkip(int bytesToPreRead, ByteBuffer inputBuf, int numBytesToSkip, boolean pushSkipToSourceStream) throws IOException {
        int expectedInpLeftAfterSkip = inputBuf.remaining() - bytesToPreRead - numBytesToSkip;
        int expectedSkippedBytes = Math.min(inputBuf.remaining() - bytesToPreRead, numBytesToSkip);

        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(inputBuf.duplicate()), supplier, 10, pushSkipToSourceStream)) {
            int cnt = 0;
            while (cnt++ < bytesToPreRead) {
                int r = is.read();
                assertNotEquals(-1, r, "Unexpected end of data.");
            }

            long res = is.skip(numBytesToSkip);
            assertEquals(expectedSkippedBytes, res);

            // verify that we are able to read rest of the input
            cnt = 0;
            while (cnt++ < expectedInpLeftAfterSkip) {
                int readRes = is.read();
                assertNotEquals(-1, readRes, "Unexpected end of data.");
            }
        }
    }

    @ParameterizedTest
    @MethodSource("provideEdgeCaseInputForMethodSkip")
    public void testEdgeCaseInputForMethodSkip(int bufferLength, long toSkip, boolean delegateSkipToSourceStream, long expected) throws IOException {
        ByteBuffer inputBuf = ByteBuffer.allocate(bufferLength);
        RANDOM.nextBytes(inputBuf.array());
        inputBuf.rewind();

        try (InputStream is = new ChunkedBytesStream(new ByteBufferInputStream(inputBuf.duplicate()), supplier, 10, delegateSkipToSourceStream)) {
            assertEquals(expected, is.skip(toSkip));
        }
    }

    private static Stream<Arguments> provideEdgeCaseInputForMethodSkip() {
        int bufferLength = 16;
        // Test toSkip larger than int and negative for both delegateToSourceStream true and false
        return Stream.of(
            Arguments.of(bufferLength, Integer.MAX_VALUE + 1L, true, bufferLength),
            Arguments.of(bufferLength, -1, true, 0),
            Arguments.of(bufferLength, Integer.MAX_VALUE + 1L, false, bufferLength),
            Arguments.of(bufferLength, -1, false, 0)
        );
    }

    private static Stream<Arguments> provideCasesWithInvalidInputsForMethodRead() {
        byte[] b = new byte[16];
        return Stream.of(
            // negative off
            Arguments.of(b, -1, b.length),
            // negative len
            Arguments.of(b, 0, -1),
            // overflow off + len
            Arguments.of(b, Integer.MAX_VALUE, 10),
            // len greater than size of target array
            Arguments.of(b, 0, b.length + 1),
            // off + len greater than size of target array
            Arguments.of(b, b.length - 1, 2)
        );
    }

    private static List<Arguments> provideSourceSkipValuesForTest() {
        ByteBuffer bufGreaterThanIntermediateBuf = ByteBuffer.allocate(16);
        RANDOM.nextBytes(bufGreaterThanIntermediateBuf.array());
        bufGreaterThanIntermediateBuf.position(bufGreaterThanIntermediateBuf.capacity());
        bufGreaterThanIntermediateBuf.flip();

        ByteBuffer bufMuchGreaterThanIntermediateBuf = ByteBuffer.allocate(100);
        RANDOM.nextBytes(bufMuchGreaterThanIntermediateBuf.array());
        bufMuchGreaterThanIntermediateBuf.position(bufMuchGreaterThanIntermediateBuf.capacity());
        bufMuchGreaterThanIntermediateBuf.flip();

        ByteBuffer emptyBuffer = ByteBuffer.allocate(2);

        ByteBuffer oneByteBuf = ByteBuffer.allocate(1).put((byte) 1);
        oneByteBuf.flip();

        List<List<Object>> testInputs = Arrays.asList(
            // empty source byte array
            Arrays.asList(0, emptyBuffer, 0),
            Arrays.asList(0, emptyBuffer, 1),
            Arrays.asList(1, emptyBuffer, 1),
            Arrays.asList(1, emptyBuffer, 0),
            // byte source array with 1 byte
            Arrays.asList(0, oneByteBuf, 0),
            Arrays.asList(0, oneByteBuf, 1),
            Arrays.asList(1, oneByteBuf, 0),
            Arrays.asList(1, oneByteBuf, 1),
            // byte source array with full read from intermediate buf
            Arrays.asList(0, bufGreaterThanIntermediateBuf.duplicate(), bufGreaterThanIntermediateBuf.capacity()),
            Arrays.asList(bufGreaterThanIntermediateBuf.capacity(), bufGreaterThanIntermediateBuf.duplicate(), 0),
            Arrays.asList(2, bufGreaterThanIntermediateBuf.duplicate(), 10),
            Arrays.asList(2, bufGreaterThanIntermediateBuf.duplicate(), 8)
        );

        Boolean[] tailArgs = new Boolean[]{true, false};
        List<Arguments> finalArguments = new ArrayList<>(2 * testInputs.size());
        for (List<Object> args : testInputs) {
            for (Boolean aBoolean : tailArgs) {
                List<Object> expandedArgs = new ArrayList<>(args);
                expandedArgs.add(aBoolean);
                finalArguments.add(Arguments.of(expandedArgs.toArray()));
            }
        }
        return finalArguments;
    }

    private static Stream<Arguments> provideSourceBytebuffersForTest() {
        ByteBuffer bufGreaterThanIntermediateBuf = ByteBuffer.allocate(16);
        RANDOM.nextBytes(bufGreaterThanIntermediateBuf.array());
        bufGreaterThanIntermediateBuf.position(bufGreaterThanIntermediateBuf.capacity());

        ByteBuffer bufMuchGreaterThanIntermediateBuf = ByteBuffer.allocate(100);
        RANDOM.nextBytes(bufMuchGreaterThanIntermediateBuf.array());
        bufMuchGreaterThanIntermediateBuf.position(bufMuchGreaterThanIntermediateBuf.capacity());

        return Stream.of(
            // empty byte array
            Arguments.of(ByteBuffer.allocate(2)),
            // byte array with 1 byte
            Arguments.of(ByteBuffer.allocate(1).put((byte) 1).flip()),
            // byte array with size < intermediate buffer
            Arguments.of(ByteBuffer.allocate(8).put("12345678".getBytes()).flip()),
            // byte array with size > intermediate buffer
            Arguments.of(bufGreaterThanIntermediateBuf.flip()),
            // byte array with size >> intermediate buffer
            Arguments.of(bufMuchGreaterThanIntermediateBuf.flip())
        );
    }
}
