/*
    * 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
    *
   *   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.exec;
  
  import java.io.IOException;
  import java.io.InputStream;
  import java.io.OutputStream;
  import java.io.PipedOutputStream;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.concurrent.Executors;
  import java.util.concurrent.ThreadFactory;
  
  import org.apache.commons.exec.util.DebugUtils;
  
  /**
   * Copies standard output and error of sub-processes to standard output and error of the parent process. If output or error stream are set to null, any feedback
   * from that stream will be lost.
   */
  public class PumpStreamHandler implements ExecuteStreamHandler {
  
      /** Three seconds timeout. */
      private static final Duration STOP_TIMEOUT_ADDITION = Duration.ofSeconds(2);
  
      /**
       * Output thread.
       */
      private Thread outputThread;
  
      /**
       * Error thread.
       */
      private Thread errorThread;
  
      /** Input thread. */
      private Thread inputThread;
  
      /** Output stream. */
      private final OutputStream outputStream;
  
      /** Error output stream. */
      private final OutputStream errorOutputStream;
  
      /** Error input stream. */
      private final InputStream inputStream;
  
      /** Pumper input stream. */
      private InputStreamPumper inputStreamPumper;
  
      /** The timeout Duration the implementation waits when stopping the pumper threads. */
      private Duration stopTimeout = Duration.ZERO;
  
      /** The last exception being caught. */
      private IOException caught;
  
      /**
       * The thread factory.
       */
      private final ThreadFactory threadFactory;
  
      /**
       * Constructs a new {@link PumpStreamHandler}.
       */
      public PumpStreamHandler() {
          this(System.out, System.err);
      }
  
      /**
       * Constructs a new {@link PumpStreamHandler}.
       *
       * @param allOutputStream the output/error {@link OutputStream}. The {@code OutputStream}
       *      implementation must be thread-safe because the output and error reader threads will
       *      concurrently write to it.
       */
      public PumpStreamHandler(final OutputStream allOutputStream) {
          this(allOutputStream, allOutputStream);
      }
  
      /**
       * Constructs a new {@link PumpStreamHandler}.
       *
       * <p>If the same {@link OutputStream} instance is used for output and error, then it must be
      * thread-safe because the output and error reader threads will concurrently write to it.
      *
      * @param outputStream      the output {@link OutputStream}.
      * @param errorOutputStream the error {@link OutputStream}.
      */
     public PumpStreamHandler(final OutputStream outputStream, final OutputStream errorOutputStream) {
         this(outputStream, errorOutputStream, null);
     }
 
     /**
      * Constructs a new {@link PumpStreamHandler}.
      *
      * <p>If the same {@link OutputStream} instance is used for output and error, then it must be
      * thread-safe because the output and error reader threads will concurrently write to it.
      *
      * @param outputStream      the output {@link OutputStream}.
      * @param errorOutputStream the error {@link OutputStream}.
      * @param inputStream       the input {@link InputStream}.
      */
     public PumpStreamHandler(final OutputStream outputStream, final OutputStream errorOutputStream, final InputStream inputStream) {
         this(Executors.defaultThreadFactory(), outputStream, errorOutputStream, inputStream);
     }
 
     /**
      * Constructs a new {@link PumpStreamHandler}.
      *
      * <p>If the same {@link OutputStream} instance is used for output and error, then it must be
      * thread-safe because the output and error reader threads will concurrently write to it.
      *
      * @param outputStream      the output {@link OutputStream}.
      * @param errorOutputStream the error {@link OutputStream}.
      * @param inputStream       the input {@link InputStream}.
      */
     private PumpStreamHandler(final ThreadFactory threadFactory, final OutputStream outputStream, final OutputStream errorOutputStream,
             final InputStream inputStream) {
         this.threadFactory = threadFactory;
         this.outputStream = outputStream;
         this.errorOutputStream = errorOutputStream;
         this.inputStream = inputStream;
     }
 
     /**
      * Create the pump to handle error output.
      *
      * @param is the {@link InputStream}.
      * @param os the {@link OutputStream}.
      */
     protected void createProcessErrorPump(final InputStream is, final OutputStream os) {
         errorThread = createPump(is, os);
     }
 
     /**
      * Create the pump to handle process output.
      *
      * @param is the {@link InputStream}.
      * @param os the {@link OutputStream}.
      */
     protected void createProcessOutputPump(final InputStream is, final OutputStream os) {
         outputThread = createPump(is, os);
     }
 
     /**
      * Creates a stream pumper to copy the given input stream to the given output stream. When the 'os' is an PipedOutputStream we are closing 'os' afterward
      * to avoid an IOException ("Write end dead").
      *
      * @param is the input stream to copy from.
      * @param os the output stream to copy into.
      * @return the stream pumper thread.
      */
     protected Thread createPump(final InputStream is, final OutputStream os) {
         return createPump(is, os, os instanceof PipedOutputStream);
     }
 
     /**
      * Creates a stream pumper to copy the given input stream to the given output stream.
      *
      * @param is                 the input stream to copy from.
      * @param os                 the output stream to copy into.
      * @param closeWhenExhausted close the output stream when the input stream is exhausted.
      * @return the stream pumper thread.
      */
     protected Thread createPump(final InputStream is, final OutputStream os, final boolean closeWhenExhausted) {
         return ThreadUtil.newThread(threadFactory, new StreamPumper(is, os, closeWhenExhausted), "CommonsExecStreamPumper-", true);
     }
 
     /**
      * Creates a stream pumper to copy the given input stream to the given output stream.
      *
      * @param is the System.in input stream to copy from.
      * @param os the output stream to copy into.
      * @return the stream pumper thread.
      */
     private Thread createSystemInPump(final InputStream is, final OutputStream os) {
         inputStreamPumper = new InputStreamPumper(is, os);
         return ThreadUtil.newThread(threadFactory, inputStreamPumper, "CommonsExecStreamPumper-", true);
     }
 
     /**
      * Gets the error stream.
      *
      * @return {@link OutputStream}.
      */
     protected OutputStream getErr() {
         return errorOutputStream;
     }
 
     /**
      * Gets the output stream.
      *
      * @return {@link OutputStream}.
      */
     protected OutputStream getOut() {
         return outputStream;
     }
 
     Duration getStopTimeout() {
         return stopTimeout;
     }
 
     /**
      * Sets the {@link InputStream} from which to read the standard error of the process.
      *
      * @param is the {@link InputStream}.
      */
     @Override
     public void setProcessErrorStream(final InputStream is) {
         if (errorOutputStream != null) {
             createProcessErrorPump(is, errorOutputStream);
         }
     }
 
     /**
      * Sets the {@link OutputStream} by means of which input can be sent to the process.
      *
      * @param os the {@link OutputStream}.
      */
     @Override
     public void setProcessInputStream(final OutputStream os) {
         if (inputStream != null) {
             if (inputStream == System.in) {
                 inputThread = createSystemInPump(inputStream, os);
             } else {
                 inputThread = createPump(inputStream, os, true);
             }
         } else {
             try {
                 os.close();
             } catch (final IOException e) {
                 final String msg = "Got exception while closing output stream";
                 DebugUtils.handleException(msg, e);
             }
         }
     }
 
     /**
      * Sets the {@link InputStream} from which to read the standard output of the process.
      *
      * @param is the {@link InputStream}.
      */
     @Override
     public void setProcessOutputStream(final InputStream is) {
         if (outputStream != null) {
             createProcessOutputPump(is, outputStream);
         }
     }
 
     /**
      * Sets maximum time to wait until output streams are exhausted when {@link #stop()} was called.
      *
      * @param timeout timeout or zero to wait forever (default).
      * @since 1.4.0
      */
     public void setStopTimeout(final Duration timeout) {
         this.stopTimeout = timeout != null ? timeout : Duration.ZERO;
     }
 
     /**
      * Sets maximum time to wait until output streams are exhausted when {@link #stop()} was called.
      *
      * @param timeout timeout in milliseconds or zero to wait forever (default).
      * @deprecated Use {@link #setStopTimeout(Duration)}.
      */
     @Deprecated
     public void setStopTimeout(final long timeout) {
         this.stopTimeout = Duration.ofMillis(timeout);
     }
 
     /**
      * Starts the {@link Thread}s.
      */
     @Override
     public void start() {
         start(outputThread);
         start(errorThread);
         start(inputThread);
     }
 
     /**
      * Starts the given {@link Thread}.
      */
     private void start(final Thread thread) {
         if (thread != null) {
             thread.start();
         }
     }
 
     /**
      * Stops pumping the streams. When a timeout is specified it is not guaranteed that the pumper threads are cleanly terminated.
      */
     @Override
     public void stop() throws IOException {
         if (inputStreamPumper != null) {
             inputStreamPumper.stopProcessing();
         }
         stop(outputThread, stopTimeout);
         stop(errorThread, stopTimeout);
         stop(inputThread, stopTimeout);
 
         if (errorOutputStream != null && errorOutputStream != outputStream) {
             try {
                 errorOutputStream.flush();
             } catch (final IOException e) {
                 final String msg = "Got exception while flushing the error stream : " + e.getMessage();
                 DebugUtils.handleException(msg, e);
             }
         }
 
         if (outputStream != null) {
             try {
                 outputStream.flush();
             } catch (final IOException e) {
                 final String msg = "Got exception while flushing the output stream";
                 DebugUtils.handleException(msg, e);
             }
         }
 
         if (caught != null) {
             throw caught;
         }
     }
 
     /**
      * Stops a pumper thread. The implementation actually waits longer than specified in 'timeout' to detect if the timeout was indeed exceeded. If the timeout
      * was exceeded an IOException is created to be thrown to the caller.
      *
      * @param thread  the thread to be stopped.
      * @param timeout the time in ms to wait to join.
      */
     private void stop(final Thread thread, final Duration timeout) {
         if (thread != null) {
             try {
                 if (timeout.equals(Duration.ZERO)) {
                     thread.join();
                 } else {
                     final Duration timeToWait = timeout.plus(STOP_TIMEOUT_ADDITION);
                     final Instant startTime = Instant.now();
                     thread.join(timeToWait.toMillis());
                     if (Instant.now().isAfter(startTime.plus(timeToWait))) {
                         caught = new ExecuteException("The stop timeout of " + timeout + " ms was exceeded", Executor.INVALID_EXITVALUE);
                     }
                 }
             } catch (final InterruptedException e) {
                 thread.interrupt();
             }
         }
     }
 
     /**
      * Stops a pumper thread. The implementation actually waits longer than specified in 'timeout' to detect if the timeout was indeed exceeded. If the timeout
      * was exceeded an IOException is created to be thrown to the caller.
      *
      * @param thread        the thread to be stopped.
      * @param timeoutMillis the time in ms to wait to join.
      */
     protected void stopThread(final Thread thread, final long timeoutMillis) {
         stop(thread, Duration.ofMillis(timeoutMillis));
     }
 }