/*
 * 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 java.io.BufferedInputStream;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;

/**
 * ChunkedBytesStream is a copy of {@link BufferedInputStream} with the following differences:
 * - Unlike {@link java.io.BufferedInputStream#skip(long)} this class could be configured to not push skip() to
 * input stream. We may want to avoid pushing this to input stream because it's implementation maybe inefficient,
 * e.g. the case of ZstdInputStream which allocates a new buffer from buffer pool, per skip call.
 * - Unlike {@link java.io.BufferedInputStream}, which allocates an intermediate buffer, this uses a buffer supplier to
 * create the intermediate buffer.
 * - Unlike {@link java.io.BufferedInputStream}, this implementation does not support {@link InputStream#mark(int)} and
 * {@link InputStream#markSupported()} will return false.
 * <p>
 * Note that:
 * - this class is not thread safe and shouldn't be used in scenarios where multiple threads access this.
 * - the implementation of this class is performance sensitive. Minor changes such as usage of ByteBuffer instead of byte[]
 * can significantly impact performance, hence, proceed with caution.
 */
public class ChunkedBytesStream extends FilterInputStream {
    /**
     * Supplies the ByteBuffer which is used as intermediate buffer to store the chunk of output data.
     */
    private final BufferSupplier bufferSupplier;
    /**
     * Intermediate buffer to store the chunk of output data. The ChunkedBytesStream is considered closed if
     * this buffer is null.
     */
    private byte[] intermediateBuf;
    /**
     * The index one greater than the index of the last valid byte in
     * the buffer.
     * This value is always in the range <code>0</code> through <code>intermediateBuf.length</code>;
     * elements <code>intermediateBuf[0]</code>  through <code>intermediateBuf[count-1]
     * </code>contain buffered input data obtained
     * from the underlying  input stream.
     */
    protected int count = 0;
    /**
     * The current position in the buffer. This is the index of the next
     * character to be read from the <code>buf</code> array.
     * <p>
     * This value is always in the range <code>0</code>
     * through <code>count</code>. If it is less
     * than <code>count</code>, then  <code>intermediateBuf[pos]</code>
     * is the next byte to be supplied as input;
     * if it is equal to <code>count</code>, then
     * the  next <code>read</code> or <code>skip</code>
     * operation will require more bytes to be
     * read from the contained  input stream.
     */
    protected int pos = 0;
    /**
     * Reference for the intermediate buffer. This reference is only kept for releasing the buffer from the
     * buffer supplier.
     */
    private final ByteBuffer intermediateBufRef;
    /**
     * If this flag is true, we will delegate the responsibility of skipping to the
     * sourceStream. This is an alternative to reading the data from source stream, storing in an intermediate buffer and
     * skipping the values using this implementation.
     */
    private final boolean delegateSkipToSourceStream;

    public ChunkedBytesStream(InputStream in, BufferSupplier bufferSupplier, int intermediateBufSize, boolean delegateSkipToSourceStream) {
        super(in);
        this.bufferSupplier = bufferSupplier;
        intermediateBufRef = bufferSupplier.get(intermediateBufSize);
        if (!intermediateBufRef.hasArray() || (intermediateBufRef.arrayOffset() != 0)) {
            throw new IllegalArgumentException("provided ByteBuffer lacks array or has non-zero arrayOffset");
        }
        intermediateBuf = intermediateBufRef.array();
        this.delegateSkipToSourceStream = delegateSkipToSourceStream;
    }

    /**
     * Check to make sure that buffer has not been nulled out due to
     * close; if not return it;
     */
    private byte[] getBufIfOpen() throws IOException {
        byte[] buffer = intermediateBuf;
        if (buffer == null)
            throw new IOException("Stream closed");
        return buffer;
    }

    /**
     * See the general contract of the <code>read</code>
     * method of <code>InputStream</code>.
     *
     * @return the next byte of data, or <code>-1</code> if the end of the
     * stream is reached.
     * @throws IOException if this input stream has been closed by
     *                     invoking its {@link #close()} method,
     *                     or an I/O error occurs.
     * @see BufferedInputStream#read()
     */
    @Override
    public int read() throws IOException {
        if (pos >= count) {
            fill();
            if (pos >= count)
                return -1;
        }

        return getBufIfOpen()[pos++] & 0xff;
    }

    /**
     * Check to make sure that underlying input stream has not been
     * nulled out due to close; if not return it;
     */
    InputStream getInIfOpen() throws IOException {
        InputStream input = in;
        if (input == null)
            throw new IOException("Stream closed");
        return input;
    }

    /**
     * Fills the intermediate buffer with more data. The amount of new data read is equal to the remaining empty space
     * in the buffer. For optimal performance, read as much data as possible in this call.
     * This method also assumes that all data has already been read in, hence pos > count.
     */
    int fill() throws IOException {
        byte[] buffer = getBufIfOpen();
        pos = 0;
        count = pos;
        int n = getInIfOpen().read(buffer, pos, buffer.length - pos);
        if (n > 0)
            count = n + pos;
        return n;
    }

    @Override
    public void close() throws IOException {
        byte[] mybuf = intermediateBuf;
        intermediateBuf = null;

        InputStream input = in;
        in = null;

        if (mybuf != null)
            bufferSupplier.release(intermediateBufRef);
        if (input != null)
            input.close();
    }

    /**
     * Reads bytes from this byte-input stream into the specified byte array,
     * starting at the given offset.
     *
     * <p> This method implements the general contract of the corresponding
     * <code>{@link InputStream#read(byte[], int, int) read}</code> method of
     * the <code>{@link InputStream}</code> class.  As an additional
     * convenience, it attempts to read as many bytes as possible by repeatedly
     * invoking the <code>read</code> method of the underlying stream.  This
     * iterated <code>read</code> continues until one of the following
     * conditions becomes true: <ul>
     *
     * <li> The specified number of bytes have been read,
     *
     * <li> The <code>read</code> method of the underlying stream returns
     * <code>-1</code>, indicating end-of-file, or
     *
     * <li> The <code>available</code> method of the underlying stream
     * returns zero, indicating that further input requests would block.
     *
     * </ul> If the first <code>read</code> on the underlying stream returns
     * <code>-1</code> to indicate end-of-file then this method returns
     * <code>-1</code>.  Otherwise this method returns the number of bytes
     * actually read.
     *
     * <p> Subclasses of this class are encouraged, but not required, to
     * attempt to read as many bytes as possible in the same fashion.
     *
     * @param b   destination buffer.
     * @param off offset at which to start storing bytes.
     * @param len maximum number of bytes to read.
     * @return the number of bytes read, or <code>-1</code> if the end of
     * the stream has been reached.
     * @throws IOException if this input stream has been closed by
     *                     invoking its {@link #close()} method,
     *                     or an I/O error occurs.
     * @see BufferedInputStream#read(byte[], int, int)
     */
    public int read(byte[] b, int off, int len) throws IOException {
        getBufIfOpen(); // Check for closed stream
        if ((off | len | (off + len) | (b.length - (off + len))) < 0) {
            throw new IndexOutOfBoundsException();
        } else if (len == 0) {
            return 0;
        }

        int n = 0;
        for (; ; ) {
            int nread = read1(b, off + n, len - n);
            if (nread <= 0)
                return (n == 0) ? nread : n;
            n += nread;
            if (n >= len)
                return n;
            // if not closed but no bytes available, return
            InputStream input = in;
            if (input != null && input.available() <= 0)
                return n;
        }
    }

    /**
     * Read characters into a portion of an array, reading from the underlying
     * stream at most once if necessary.
     * <p>
     * Note - Implementation copied from {@link BufferedInputStream}. Slight modification done to remove
     * the mark position.
     */
    private int read1(byte[] b, int off, int len) throws IOException {
        int avail = count - pos;
        if (avail <= 0) {
            /* If the requested length is at least as large as the buffer, and
               if there is no mark/reset activity, do not bother to copy the
               bytes into the local buffer.  In this way buffered streams will
               cascade harmlessly. */
            if (len >= getBufIfOpen().length) {
                return getInIfOpen().read(b, off, len);
            }
            fill();
            avail = count - pos;
            if (avail <= 0) return -1;
        }
        int cnt = Math.min(avail, len);
        System.arraycopy(getBufIfOpen(), pos, b, off, cnt);
        pos += cnt;
        return cnt;
    }

    /**
     * Skips over and discards exactly {@code n} bytes of data from this input stream.
     * If {@code n} is zero, then no bytes are skipped.
     * If {@code n} is negative, then no bytes are skipped.
     * <p>
     * This method blocks until the requested number of bytes has been skipped, end of file is reached, or an
     * exception is thrown.
     * <p> If end of stream is reached before the stream is at the desired position, then the bytes skipped till than pointed
     * are returned.
     * <p> If an I/O error occurs, then the input stream may be in an inconsistent state. It is strongly recommended that the
     * stream be promptly closed if an I/O error occurs.
     * <p>
     * This method first skips and discards bytes in the intermediate buffer.
     * After that, depending on the value of {@link #delegateSkipToSourceStream}, it either delegates the skipping of
     * bytes to the sourceStream or it reads the data from input stream in chunks, copies the data into intermediate
     * buffer and skips it.
     * <p>
     * Starting JDK 12, a new method was introduced in InputStream, skipNBytes which has a similar behaviour as
     * this method.
     *
     * @param toSkip the number of bytes to be skipped.
     * @return the actual number of bytes skipped which might be zero.
     * @throws IOException if this input stream has been closed by invoking its {@link #close()} method,
     *                     {@code in.skip(n)} throws an IOException, or an I/O error occurs.
     */
    @Override
    public long skip(long toSkip) throws IOException {
        getBufIfOpen(); // Check for closed stream
        if (toSkip <= 0) {
            return 0;
        }

        long remaining = toSkip;

        // Skip bytes stored in intermediate buffer first
        int avail = count - pos;
        int bytesSkipped = (int) Math.min(avail, remaining);
        pos += bytesSkipped;
        remaining -= bytesSkipped;

        while (remaining > 0) {
            if (delegateSkipToSourceStream) {
                // Use sourceStream's skip() to skip the rest.
                // conversion to int is acceptable because toSkip and remaining are int.
                long delegateBytesSkipped = getInIfOpen().skip(remaining);
                if (delegateBytesSkipped == 0) {
                    // read one byte to check for EOS
                    if (read() == -1) {
                        break;
                    }
                    // one byte read so decrement number to skip
                    remaining--;
                } else if (delegateBytesSkipped > remaining || delegateBytesSkipped < 0) { // skipped negative or too many bytes
                    throw new IOException("Unable to skip exactly");

                }
                remaining -= delegateBytesSkipped;
            } else {
                // skip from intermediate buffer, filling it first (if required)
                if (pos >= count) {
                    fill();
                    // if we don't have data in intermediate buffer after fill, then stop skipping
                    if (pos >= count)
                        break;
                }
                avail = count - pos;
                bytesSkipped = (int) Math.min(avail, remaining);
                pos += bytesSkipped;
                remaining -= bytesSkipped;
            }
        }
        return toSkip - remaining;
    }

    // visible for testing
    public InputStream sourceStream() {
        return in;
    }

    /**
     * Returns an estimate of the number of bytes that can be read (or
     * skipped over) from this input stream without blocking by the next
     * invocation of a method for this input stream. The next invocation might be
     * the same thread or another thread.  A single read or skip of this
     * many bytes will not block, but may read or skip fewer bytes.
     * <p>
     * This method returns the sum of the number of bytes remaining to be read in
     * the buffer (<code>count&nbsp;- pos</code>) and the result of calling the
     * {@link java.io.FilterInputStream#in in}.available().
     *
     * @return an estimate of the number of bytes that can be read (or skipped
     * over) from this input stream without blocking.
     * @throws IOException if this input stream has been closed by
     *                     invoking its {@link #close()} method,
     *                     or an I/O error occurs.
     * @see BufferedInputStream#available()
     */
    @Override
    public synchronized int available() throws IOException {
        int n = count - pos;
        int avail = getInIfOpen().available();
        return n > (Integer.MAX_VALUE - avail)
            ? Integer.MAX_VALUE
            : n + avail;
    }
}
