/*
    * Licensed 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
    *
    *     https://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.commons.io.input;
  
  import static org.apache.commons.io.IOUtils.EOF;
  
  // import javax.annotation.concurrent.GuardedBy;
  import java.io.EOFException;
  import java.io.FilterInputStream;
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.InterruptedIOException;
  import java.nio.ByteBuffer;
  import java.util.Objects;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.locks.Condition;
  import java.util.concurrent.locks.ReentrantLock;
  
  import org.apache.commons.io.IOUtils;
  import org.apache.commons.io.build.AbstractStreamBuilder;
  
  /**
   * Implements {@link InputStream} to asynchronously read ahead from an underlying input stream when a specified amount of data has been read from the current
   * buffer. It does so by maintaining two buffers: an active buffer and a read ahead buffer. The active buffer contains data which should be returned when a
   * read() call is issued. The read ahead buffer is used to asynchronously read from the underlying input stream. When the current active buffer is exhausted, we
   * flip the two buffers so that we can start reading from the read ahead buffer without being blocked by disk I/O.
   * <p>
   * To build an instance, use {@link Builder}.
   * </p>
   * <p>
   * This class was ported and adapted from Apache Spark commit 933dc6cb7b3de1d8ccaf73d124d6eb95b947ed19.
   * </p>
   *
   * @see Builder
   * @since 2.9.0
   */
  public class ReadAheadInputStream extends FilterInputStream {
  
      // @formatter:off
      /**
       * Builds a new {@link ReadAheadInputStream}.
       *
       * <p>
       * For example:
       * </p>
       * <pre>{@code
       * ReadAheadInputStream s = ReadAheadInputStream.builder()
       *   .setPath(path)
       *   .setExecutorService(Executors.newSingleThreadExecutor(ReadAheadInputStream::newThread))
       *   .get();}
       * </pre>
       *
       * @see #get()
       * @since 2.12.0
       */
      // @formatter:on
      public static class Builder extends AbstractStreamBuilder<ReadAheadInputStream, Builder> {
  
          private ExecutorService executorService;
  
          /**
           * Constructs a new builder of {@link ReadAheadInputStream}.
           */
          public Builder() {
              // empty
          }
  
          /**
           * Builds a new {@link ReadAheadInputStream}.
           * <p>
           * You must set an aspect that supports {@link #getInputStream()}, otherwise, this method throws an exception.
           * </p>
           * <p>
           * This builder uses the following aspects:
           * </p>
           * <ul>
           * <li>{@link #getInputStream()} gets the target aspect.</li>
           * <li>{@link #getBufferSize()}</li>
           * <li>{@link ExecutorService}</li>
           * </ul>
           *
           * @return a new instance.
           * @throws IllegalStateException         if the {@code origin} is {@code null}.
           * @throws UnsupportedOperationException if the origin cannot be converted to an {@link InputStream}.
           * @throws IOException                   if an I/O error occurs converting to an {@link InputStream} using {@link #getInputStream()}.
          * @see #getInputStream()
          * @see #getBufferSize()
          * @see #getUnchecked()
          */
         @Override
         public ReadAheadInputStream get() throws IOException {
             return new ReadAheadInputStream(this);
         }
 
         /**
          * Sets the executor service for the read-ahead thread.
          *
          * @param executorService the executor service for the read-ahead thread.
          * @return {@code this} instance.
          */
         public Builder setExecutorService(final ExecutorService executorService) {
             this.executorService = executorService;
             return this;
         }
 
     }
 
     private static final ThreadLocal<byte[]> BYTE_ARRAY_1 = ThreadLocal.withInitial(() -> new byte[1]);
 
     /**
      * Constructs a new {@link Builder}.
      *
      * @return a new {@link Builder}.
      * @since 2.12.0
      */
     public static Builder builder() {
         return new Builder();
     }
 
     /**
      * Constructs a new daemon thread.
      *
      * @param r the thread's runnable.
      * @return a new daemon thread.
      */
     private static Thread newDaemonThread(final Runnable r) {
         final Thread thread = new Thread(r, "commons-io-read-ahead");
         thread.setDaemon(true);
         return thread;
     }
 
     /**
      * Constructs a new daemon executor service.
      *
      * @return a new daemon executor service.
      */
     private static ExecutorService newExecutorService() {
         return Executors.newSingleThreadExecutor(ReadAheadInputStream::newDaemonThread);
     }
 
     private final ReentrantLock stateChangeLock = new ReentrantLock();
 
     // @GuardedBy("stateChangeLock")
     private ByteBuffer activeBuffer;
 
     // @GuardedBy("stateChangeLock")
     private ByteBuffer readAheadBuffer;
 
     // @GuardedBy("stateChangeLock")
     private boolean endOfStream;
 
     // @GuardedBy("stateChangeLock")
     // true if async read is in progress
     private boolean readInProgress;
 
     // @GuardedBy("stateChangeLock")
     // true if read is aborted due to an exception in reading from underlying input stream.
     private boolean readAborted;
 
     // @GuardedBy("stateChangeLock")
     private Throwable readException;
 
     // @GuardedBy("stateChangeLock")
     // whether the close method is called.
     private boolean isClosed;
 
     // @GuardedBy("stateChangeLock")
     // true when the close method will close the underlying input stream. This is valid only if
     // `isClosed` is true.
     private boolean isUnderlyingInputStreamBeingClosed;
 
     // @GuardedBy("stateChangeLock")
     // whether there is a read ahead task running,
     private boolean isReading;
 
     // Whether there is a reader waiting for data.
     private final AtomicBoolean isWaiting = new AtomicBoolean();
 
     private final ExecutorService executorService;
 
     private final boolean shutdownExecutorService;
 
     private final Condition asyncReadComplete = stateChangeLock.newCondition();
 
     @SuppressWarnings("resource")
     private ReadAheadInputStream(final Builder builder) throws IOException {
         this(builder.getInputStream(), builder.getBufferSize(), builder.executorService != null ? builder.executorService : newExecutorService(),
                 builder.executorService == null);
     }
 
     /**
      * Constructs an instance with the specified buffer size and read-ahead threshold
      *
      * @param inputStream       The underlying input stream.
      * @param bufferSizeInBytes The buffer size.
      * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
      */
     @Deprecated
     public ReadAheadInputStream(final InputStream inputStream, final int bufferSizeInBytes) {
         this(inputStream, bufferSizeInBytes, newExecutorService(), true);
     }
 
     /**
      * Constructs an instance with the specified buffer size and read-ahead threshold
      *
      * @param inputStream       The underlying input stream.
      * @param bufferSizeInBytes The buffer size.
      * @param executorService   An executor service for the read-ahead thread.
      * @deprecated Use {@link #builder()}, {@link Builder}, and {@link Builder#get()}
      */
     @Deprecated
     public ReadAheadInputStream(final InputStream inputStream, final int bufferSizeInBytes, final ExecutorService executorService) {
         this(inputStream, bufferSizeInBytes, executorService, false);
     }
 
     /**
      * Constructs an instance with the specified buffer size and read-ahead threshold
      *
      * @param inputStream             The underlying input stream.
      * @param bufferSizeInBytes       The buffer size.
      * @param executorService         An executor service for the read-ahead thread.
      * @param shutdownExecutorService Whether or not to shut down the given ExecutorService on close.
      */
     private ReadAheadInputStream(final InputStream inputStream, final int bufferSizeInBytes, final ExecutorService executorService,
             final boolean shutdownExecutorService) {
         super(Objects.requireNonNull(inputStream, "inputStream"));
         if (bufferSizeInBytes <= 0) {
             throw new IllegalArgumentException(String.format("bufferSizeInBytes <= 0, bufferSizeInBytes = %,d", bufferSizeInBytes));
         }
         this.executorService = Objects.requireNonNull(executorService, "executorService");
         this.shutdownExecutorService = shutdownExecutorService;
         this.activeBuffer = ByteBuffer.allocate(bufferSizeInBytes);
         this.readAheadBuffer = ByteBuffer.allocate(bufferSizeInBytes);
         this.activeBuffer.flip();
         this.readAheadBuffer.flip();
     }
 
     @Override
     public int available() throws IOException {
         stateChangeLock.lock();
         // Make sure we have no integer overflow.
         try {
             return (int) Math.min(Integer.MAX_VALUE, (long) activeBuffer.remaining() + readAheadBuffer.remaining());
         } finally {
             stateChangeLock.unlock();
         }
     }
 
     private void checkReadException() throws IOException {
         if (readAborted) {
             if (readException instanceof IOException) {
                 throw (IOException) readException;
             }
             throw new IOException(readException);
         }
     }
 
     @Override
     public void close() throws IOException {
         boolean isSafeToCloseUnderlyingInputStream = false;
         stateChangeLock.lock();
         try {
             if (isClosed) {
                 return;
             }
             isClosed = true;
             if (!isReading) {
                 // Nobody is reading, so we can close the underlying input stream in this method.
                 isSafeToCloseUnderlyingInputStream = true;
                 // Flip this to make sure the read ahead task will not close the underlying input stream.
                 isUnderlyingInputStreamBeingClosed = true;
             }
         } finally {
             stateChangeLock.unlock();
         }
 
         if (shutdownExecutorService) {
             try {
                 executorService.shutdownNow();
                 executorService.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS);
             } catch (final InterruptedException e) {
                 final InterruptedIOException iio = new InterruptedIOException(e.getMessage());
                 iio.initCause(e);
                 throw iio;
             } finally {
                 if (isSafeToCloseUnderlyingInputStream) {
                     super.close();
                 }
             }
         }
     }
 
     private void closeUnderlyingInputStreamIfNecessary() {
         boolean needToCloseUnderlyingInputStream = false;
         stateChangeLock.lock();
         try {
             isReading = false;
             if (isClosed && !isUnderlyingInputStreamBeingClosed) {
                 // close method cannot close underlyingInputStream because we were reading.
                 needToCloseUnderlyingInputStream = true;
             }
         } finally {
             stateChangeLock.unlock();
         }
         if (needToCloseUnderlyingInputStream) {
             try {
                 super.close();
             } catch (final IOException ignored) {
                 // TODO Rethrow as UncheckedIOException?
             }
         }
     }
 
     private boolean isEndOfStream() {
         return !activeBuffer.hasRemaining() && !readAheadBuffer.hasRemaining() && endOfStream;
     }
 
     @Override
     public int read() throws IOException {
         if (activeBuffer.hasRemaining()) {
             // short path - just get one byte.
             return activeBuffer.get() & 0xFF;
         }
         final byte[] oneByteArray = BYTE_ARRAY_1.get();
         oneByteArray[0] = 0;
         return read(oneByteArray, 0, 1) == EOF ? EOF : oneByteArray[0] & 0xFF;
     }
 
     @Override
     public int read(final byte[] b, final int offset, int len) throws IOException {
         IOUtils.checkFromIndexSize(b, offset, len);
         if (len == 0) {
             return 0;
         }
 
         if (!activeBuffer.hasRemaining()) {
             // No remaining in active buffer - lock and switch to write ahead buffer.
             stateChangeLock.lock();
             try {
                 waitForAsyncReadComplete();
                 if (!readAheadBuffer.hasRemaining()) {
                     // The first read.
                     readAsync();
                     waitForAsyncReadComplete();
                     if (isEndOfStream()) {
                         return EOF;
                     }
                 }
                 // Swap the newly read ahead buffer in place of empty active buffer.
                 swapBuffers();
                 // After swapping buffers, trigger another async read for read ahead buffer.
                 readAsync();
             } finally {
                 stateChangeLock.unlock();
             }
         }
         len = Math.min(len, activeBuffer.remaining());
         activeBuffer.get(b, offset, len);
 
         return len;
     }
 
     /**
      * Reads data from underlyingInputStream to readAheadBuffer asynchronously.
      *
      * @throws IOException if an I/O error occurs.
      */
     private void readAsync() throws IOException {
         stateChangeLock.lock();
         final byte[] arr;
         try {
             arr = readAheadBuffer.array();
             if (endOfStream || readInProgress) {
                 return;
             }
             checkReadException();
             readAheadBuffer.position(0);
             readAheadBuffer.flip();
             readInProgress = true;
         } finally {
             stateChangeLock.unlock();
         }
         executorService.execute(() -> {
             stateChangeLock.lock();
             try {
                 if (isClosed) {
                     readInProgress = false;
                     return;
                 }
                 // Flip this so that the close method will not close the underlying input stream when we
                 // are reading.
                 isReading = true;
             } finally {
                 stateChangeLock.unlock();
             }
 
             // Please note that it is safe to release the lock and read into the read ahead buffer
             // because either of following two conditions will hold:
             //
             // 1. The active buffer has data available to read so the reader will not read from the read ahead buffer.
             //
             // 2. This is the first time read is called or the active buffer is exhausted, in that case the reader waits
             // for this async read to complete.
             //
             // So there is no race condition in both the situations.
             int read = 0;
             int off = 0;
             int len = arr.length;
             Throwable exception = null;
             try {
                 // try to fill the read ahead buffer.
                 // if a reader is waiting, possibly return early.
                 do {
                     read = in.read(arr, off, len);
                     if (read <= 0) {
                         break;
                     }
                     off += read;
                     len -= read;
                 } while (len > 0 && !isWaiting.get());
             } catch (final Throwable ex) {
                 exception = ex;
                 if (ex instanceof Error) {
                     // `readException` may not be reported to the user. Rethrow Error to make sure at least
                     // The user can see Error in UncaughtExceptionHandler.
                     throw (Error) ex;
                 }
             } finally {
                 stateChangeLock.lock();
                 try {
                     readAheadBuffer.limit(off);
                     if (read < 0 || exception instanceof EOFException) {
                         endOfStream = true;
                     } else if (exception != null) {
                         readAborted = true;
                         readException = exception;
                     }
                     readInProgress = false;
                     signalAsyncReadComplete();
                 } finally {
                     stateChangeLock.unlock();
                 }
                 closeUnderlyingInputStreamIfNecessary();
             }
         });
     }
 
     private void signalAsyncReadComplete() {
         stateChangeLock.lock();
         try {
             asyncReadComplete.signalAll();
         } finally {
             stateChangeLock.unlock();
         }
     }
 
     @Override
     public long skip(final long n) throws IOException {
         if (n <= 0L) {
             return 0L;
         }
         if (n <= activeBuffer.remaining()) {
             // Only skipping from active buffer is sufficient
             activeBuffer.position((int) n + activeBuffer.position());
             return n;
         }
         stateChangeLock.lock();
         final long skipped;
         try {
             skipped = skipInternal(n);
         } finally {
             stateChangeLock.unlock();
         }
         return skipped;
     }
 
     /**
      * Internal skip function which should be called only from skip(). The assumption is that the stateChangeLock is already acquired in the caller before
      * calling this function.
      *
      * @param n the number of bytes to be skipped.
      * @return the actual number of bytes skipped.
      * @throws IOException if an I/O error occurs.
      */
     private long skipInternal(final long n) throws IOException {
         if (!stateChangeLock.isLocked()) {
             throw new IllegalStateException("Expected stateChangeLock to be locked");
         }
         waitForAsyncReadComplete();
         if (isEndOfStream()) {
             return 0;
         }
         if (available() >= n) {
             // we can skip from the internal buffers
             int toSkip = (int) n;
             // We need to skip from both active buffer and read ahead buffer
             toSkip -= activeBuffer.remaining();
             if (toSkip <= 0) { // skipping from activeBuffer already handled.
                 throw new IllegalStateException("Expected toSkip > 0, actual: " + toSkip);
             }
             activeBuffer.position(0);
             activeBuffer.flip();
             readAheadBuffer.position(toSkip + readAheadBuffer.position());
             swapBuffers();
             // Trigger async read to emptied read ahead buffer.
             readAsync();
             return n;
         }
         final int skippedBytes = available();
         final long toSkip = n - skippedBytes;
         activeBuffer.position(0);
         activeBuffer.flip();
         readAheadBuffer.position(0);
         readAheadBuffer.flip();
         final long skippedFromInputStream = in.skip(toSkip);
         readAsync();
         return skippedBytes + skippedFromInputStream;
     }
 
     /**
      * Flips the active and read ahead buffer
      */
     private void swapBuffers() {
         final ByteBuffer temp = activeBuffer;
         activeBuffer = readAheadBuffer;
         readAheadBuffer = temp;
     }
 
     private void waitForAsyncReadComplete() throws IOException {
         stateChangeLock.lock();
         try {
             isWaiting.set(true);
             // There is only one reader, and one writer, so the writer should signal only once,
             // but a while loop checking the wake-up condition is still needed to avoid spurious wakeups.
             while (readInProgress) {
                 asyncReadComplete.await();
             }
         } catch (final InterruptedException e) {
             final InterruptedIOException iio = new InterruptedIOException(e.getMessage());
             iio.initCause(e);
             throw iio;
         } finally {
             try {
                 isWaiting.set(false);
             } finally {
                 stateChangeLock.unlock();
             }
         }
         checkReadException();
     }
 }