package com.atlassian.bitbucket.mesh.process.nu;

import com.atlassian.bitbucket.mesh.process.*;
import com.google.common.math.LongMath;
import com.zaxxer.nuprocess.NuProcess;
import com.zaxxer.nuprocess.NuProcessHandler;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;
import java.time.Duration;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BooleanSupplier;

import static java.util.Objects.requireNonNull;

class StdioNuProcessHandler<T> implements NuProcessHandler {

    private static final Logger log = LoggerFactory.getLogger(StdioNuProcessHandler.class);
    //A logger for _process-level_ debugging. This should not be used for _handler-level_ logging,
    //like errors from the callback or similar. This logger is for journaling process start/exit.
    private static final Logger processLog = LoggerFactory.getLogger(OsProcess.class);

    private final String commandLine;
    private final ProcessCompletionHandler completionHandler;
    private final long executionTimeout;
    private final CompletableFuture<T> future;
    private final Duration idleInterval;
    private final AtomicReference<HandlerState> state;
    private final ProcessStdioHandler<T> stdioHandler;
    private final ScheduledExecutorService timeoutExecutor;
    private final Path workDir;

    private NuOsProcess process;
    private long startTimestamp;
    private StringBuilder stderrBuffer;
    private long stderrRead;
    private long stdinWritten;
    private long stdoutRead;
    private Throwable thrown;

    //Timeouts are applied by an executor service thread, so these fields will be read/written
    //by multiple threads and must be volatile
    private volatile long idleTimeout;
    private volatile ScheduledFuture<?> timeoutFuture;

    StdioNuProcessHandler(ProcessParameters<T> parameters, ScheduledExecutorService timeoutExecutor) {
        this.timeoutExecutor = timeoutExecutor;

        commandLine = requireNonNull(parameters, "parameters").toString();
        completionHandler = parameters.getCompletionHandler();
        future = new CompletableFuture<>();
        state = new AtomicReference<>(HandlerState.CREATED);
        stdioHandler = parameters.getStdioHandler();
        workDir = parameters.getWorkDir();

        //Prepare the execution timeout. If the duration is unlimited, apply no execution timeout.
        //Otherwise, add the specified interval to the current time to calculate the timeout.
        Duration execution = parameters.getExecutionInterval();
        executionTimeout = isUnlimited(execution) ?
                Long.MAX_VALUE :
                LongMath.saturatedAdd(System.currentTimeMillis(), execution.toMillis());

        //Prepare the initial idle timeout. If the duration is unlimited, default to 60 seconds so
        //processes always have an idle timeout.
        Duration idle = parameters.getIdleInterval();
        idleInterval = isUnlimited(idle) ? Duration.ofSeconds(60L) : idle;
        idleTimeout = LongMath.saturatedAdd(System.currentTimeMillis(), idleInterval.toMillis());
    }

    @Override
    public void onExit(int exitCode) {
        if (processLog.isDebugEnabled()) {
            processLog.debug("{}: [{}] exited {} in {}ms (stdin: {}, stdout: {}, stderr: {})",
                    process.getPid(), commandLine, exitCode, System.currentTimeMillis() - startTimestamp,
                    stdinWritten, stdoutRead, stderrRead);
        }

        ScheduledFuture<?> timeoutFuture = this.timeoutFuture;
        if (!(timeoutFuture == null || timeoutFuture.isDone())) {
            //If we scheduled a task to apply our timeout, and it hasn't already completed, cancel it
            timeoutFuture.cancel(false);
        }

        try {
            stdioHandler.onExit(exitCode);
        } catch (Throwable t) {
            if (thrown != null) {
                t.addSuppressed(thrown);
            }
            thrown = t;
        }

        finish(exitCode);
    }

    @Override
    public void onPreStart(NuProcess nuProcess) {
        process = new NuOsProcess(nuProcess, commandLine);

        invokeCallback(() -> stdioHandler.onPreStart(process));
    }

    @Override
    public void onStart(NuProcess nuProcess) {
        maybeMarkStarted(); //If onStart is called, mark the process as started
    }

    @Override
    public void onStderr(ByteBuffer buffer, boolean closed) {
        if (buffer.hasRemaining()) {
            maybeMarkStarted();  //If we receive stderr, mark the process as started
        }

        //Remember the buffer's initial position. Don't assume 0, because multibyte characters may be
        //split between onStderr invocations and, when that happens, 1-3 partial bytes could still be
        //in the buffer from the previous callback
        int initialPosition = buffer.position();
        int initialRemaining = buffer.remaining(); //For calculating stderrRead

        invokeCallback(() -> {
            //Unlike onStdout, this doesn't discard the buffer if there's a failure. Instead, the buffer
            //is left intact so it can be appended to stderrBuffer below
            if (!maybeDestroy()) {
                stdioHandler.onStderr(buffer, closed);
                resetIdleTimeout();
            }
        });

        if (buffer.hasRemaining() && buffer.position() == initialPosition) {
            //If the buffer isn't empty and the delegate handler didn't consume any bytes, assume they
            //don't want to and buffer stderr to pass to the completion handler
            if (stderrBuffer == null) {
                stderrBuffer = new StringBuilder();
            }
            //TODO: If stderr is split across multiple ByteBuffers, does this work correctly?
            stderrBuffer.append(StandardCharsets.UTF_8.decode(buffer));
        }

        //This is done here, rather than after (or in) handleOutput because we want to count stderr
        //bytes that were added to stderrBuffer rather than consumed by the handler
        stderrRead += (initialRemaining - buffer.remaining());
    }

    @Override
    public boolean onStdinReady(ByteBuffer buffer) {
        maybeMarkStarted(); //If stdin is requested, mark the process as started

        return invokeCallback(() -> {
            if (maybeDestroy()) {
                return false;
            }

            boolean result = stdioHandler.onStdinReady(buffer);
            resetIdleTimeout();

            //The handler is required to flip the buffer, so whatever data now remains in the buffer
            //is how many bytes the handler wrote to stdin
            stdinWritten += buffer.remaining();
            return result;
        });
    }

    @Override
    public void onStdout(ByteBuffer buffer, boolean closed) {
        if (buffer.hasRemaining()) {
            maybeMarkStarted(); //If we receive stdout, mark the process as started
        }

        BooleanSupplier callback = () -> {
            if (maybeDestroy()) {
                //If processing is "done" (generally either meaning it has failed, timed out or been
                //canceled), don't pump any more stdout to the callback
                return false;
            }

            //Otherwise, invoke the handler to process the buffer. If this throws, invokeCallback will
            //will return false below and the buffer will be discarded
            stdioHandler.onStdout(buffer, closed);
            resetIdleTimeout();

            //NuProcess calls compact() after onStdout returns, and then checks hasRemaining() to ensure
            //the buffer isn't full. If it is, NuProcess abandons the process entirely, no longer pumping
            //I/O or actively checking for termination, and the process is leaked as a zombie.
            //
            //To prevent that, if the buffer is full after calling NioStdioHandler we throw an exception.
            //This will terminate the the process and, below, discard the stdout buffer. This is only done
            //if stdout isn't closed because NuProcess doesn't abort if the buffer is full after the final
            //onStdout callback (since the buffer will not be used again)
            if (!closed && buffer.remaining() == buffer.capacity()) {
                log.warn("{}: {}.onStdout has not consumed any data. Aborting [{}]",
                        process.getPid(), stdioHandler.getClass().getSimpleName(), commandLine);
                throw new IllegalStateException("Process aborted due to unconsumed output");
            }

            return true;
        };

        int initialRemaining = buffer.remaining();
        if (invokeCallback(callback)) {
            stdoutRead += (initialRemaining - buffer.remaining());
        } else if (buffer.hasRemaining()) {
            //If processing has already completed (maybeDestroy() returned true) or failed (onStdout
            //threw an exception), clear the buffer to ensure NuProcess doesn't discard the process
            log.debug("{}: Discarding {} bytes of stdout", process.getPid(), buffer.remaining());
            buffer.position(buffer.position() + buffer.remaining());
        }
    }

    /**
     * Returns a {@link Future}, which can be used to wait for the process to complete asynchronously
     * and retrieve its result, or cancel the process.
     *
     * @return a future result for the process
     */
    @Nonnull
    Future<T> asFuture() {
        if (state.get() == HandlerState.CREATED && !future.isDone()) {
            //If the process could not be started at the operating system level NuProcess _attempts_ to
            //ensure onExit is called, which should have completed the future. If the future isn't done,
            //assume the onExit call was missed/evaded somehow and complete the future here.
            //
            //Note: The exit code here isn't important (i.e. it's not a "magic number"); it's just a
            //      quasi-distinctive arbitrary value
            finish(-1001);
        }

        future.whenComplete((result, exception) -> {
            //If the Future is canceled, the process may still be running. This attempts to ensure it's
            //stopped. Other failures, like command timeouts, could also leave a running process behind
            process.terminateIfRunning();
        });

        //Warning: Don't return the CompletableFuture from whenComplete(...). That's added as a "listener",
        //to be invoked no matter what. Since whenComplete returns a new future, if we return that things
        //like cancel(boolean) don't work because they cancel the "wrong" future
        return future;
    }

    /**
     * Retrieves the result for an <i>already-completed</i> process.
     *
     * @return the process's result
     */
    @Nullable
    T asResult() {
        if (!future.isDone()) {
            //It is expected at this point that the process has already completed. If no result is
            //available on the future, set one. (Most likely the process was not actually started)
            //
            //Note: The exit code here isn't important (i.e. it's not a "magic number"); it's just a
            //      quasi-distinctive arbitrary value
            finish(-1002);
        }

        //finish(int) will always complete the future, so getNow here will only return null if the
        //ProcessStdioHandler returned null
        try {
            return future.getNow(null);
        } catch (CompletionException e) {
            throw (RuntimeException) e.getCause();
        }
    }

    private static boolean isUnlimited(Duration duration) {
        return duration == null || duration.isNegative() || duration.isZero();
    }

    private void callCompletionHandler(int exitCode) {
        ProcessSummary summary = new ProcessSummary.Builder(commandLine)
                .canceled(isCanceled())
                .exitCode(exitCode)
                .stderr(stderrBuffer == null ? null : StringUtils.chomp(stderrBuffer.toString()))
                .stderrRead(stderrRead)
                .stdinWritten(stdinWritten)
                .stdoutRead(stdoutRead)
                .timedOut(isTimedOut())
                .thrown(thrown)
                .build();
        completionHandler.onComplete(summary);
    }

    private void ensureStarted() {
        if (state.get() == HandlerState.CREATED) {
            IllegalStateException e = new IllegalStateException(
                    "[" + commandLine + "] could not be started", null);
            if (thrown == null) { //onPreStart could have thrown and set an exception already
                thrown = e;
            } else {
                thrown.addSuppressed(e);
            }
        }
    }

    private void finish(int exitCode) {
        if (state.get() == HandlerState.FINISHED) {
            return;
        }
        ensureStarted();

        try {
            callCompletionHandler(exitCode);

            //If we make it here, the completion handler didn't throw, so we're good to complete the future
            //with the I/O handler's result. (Note that, if the future was canceled, this will do nothing)
            future.complete(stdioHandler.getResult());
        } catch (RuntimeException e) {
            //The completion handler reviews any previously-thrown exception, as well as the exit code, so
            //any exception it throws takes precedence over any previous exception
            future.completeExceptionally(e);
        } finally {
            state.set(HandlerState.FINISHED); //Intentional unconditional assignment to FINISHED
        }
    }

    /**
     * @return the delay between now and the "nearest" timeout, which may be zero or negative if the timeout
     *         has already expired
     */
    private long getTimeoutDelay() {
        return Math.min(executionTimeout, idleTimeout) - System.currentTimeMillis();
    }

    private void invokeCallback(Runnable callback) {
        invokeCallback(() -> {
            callback.run();

            return true;
        });
    }

    private boolean invokeCallback(BooleanSupplier callback) {
        try {
            return callback.getAsBoolean();
        } catch (Throwable t) { //Yes, Throwable. We need to catch Errors too in this very special case
            if (thrown == null) {
                thrown = t;
            } else {
                thrown.addSuppressed(t);
            }

            //If one of our handlers fails and the process is still running, terminate it
            if (log.isDebugEnabled()) {
                log.debug("{}: Destroying [{}] after handler failure", process.getPid(), commandLine, t);
            }
            process.destroyIfRunning();
        }

        return false;
    }

    private boolean isCanceled() {
        return future.isCancelled() || process.isCanceled();
    }

    private boolean isTimedOut() {
        long now = System.currentTimeMillis();

        return now > executionTimeout || now > idleTimeout;
    }

    /**
     * {@link NuOsProcess#destroyIfRunning() Destroys} the process if it has been canceled or timed out, or if
     * an exception has been thrown by one of the handler methods.
     *
     * @return {@code true} if the process was destroyed; otherwise, {@code false} if the process was not
     *         destroyed (so subsequent callbacks should still be run)
     */
    private boolean maybeDestroy() {
        boolean canceled = isCanceled();
        boolean timedOut = isTimedOut();
        if (thrown != null || canceled || timedOut) {
            //Each time a callback is invoked after a failure, attempt to destroy the process.
            //The goal here is twofold:
            //- Bring the process down so NuProcess will stop pumping stderr, stdin and stdout
            //- Reach onExit so we can wrap up
            if (log.isDebugEnabled()) {
                log.debug("{}: Dropping [{}] callback invocation (Canceled: {}; Timed out: {})",
                        process.getPid(), commandLine, canceled, timedOut, thrown);
            }
            process.destroyIfRunning();

            return true;
        }

        return false;
    }

    private void maybeMarkStarted() {
        if (state.compareAndSet(HandlerState.CREATED, HandlerState.STARTED)) {
            if (processLog.isTraceEnabled()) {
                processLog.trace("{}: [{}] started (cwd: {})", process.getPid(), commandLine, workDir);
            }
            startTimestamp = System.currentTimeMillis();

            //When we first flip the state from CREATED to STARTED, invoke the delegate's onStart callback
            invokeCallback(() -> stdioHandler.onStart(process));

            //This is called after invoking the delegate handler to give its onStart an opportunity to cancel
            //the process or throw an exception, either of which would obviate the need to schedule a timeout
            maybeScheduleTimeout();
        }
    }

    private void maybeScheduleTimeout() {
        if (thrown == null && !isCanceled()) {
            //When scheduling the _initial_ timeout check, we want to ensure that, even if the caller
            //has specified a nonsensically short timeout, we give the process a small amount of time
            //to make progress. Otherwise we could end up trying to destroy the process while we're still
            //in onStart (since the timeout is applied on an executor thread), which can result in zombie
            //processes on Linux/macOS. Subsequent timeout checks do not apply this minimum delay
            long delay = Math.max(getTimeoutDelay(), 5000L);

            timeoutFuture = timeoutExecutor.schedule(new TimeoutTask(), delay, TimeUnit.MILLISECONDS);
        }
    }

    private void resetIdleTimeout() {
        //An idle interval will always be set, either because the caller provided one or because we
        //applied a default, so we can unconditionally update
        idleTimeout = System.currentTimeMillis() + idleInterval.toMillis();
    }

    /**
     * Enumerates some relevant states for the {@code NioProcessHandler}.
     */
    private enum HandlerState {

        /**
         * Indicates the handler has been created, but the process has not yet been started.
         */
        CREATED,
        /**
         * Indicates the process has finished (successfully or not).
         */
        FINISHED,
        /**
         * Indicates the process was successfully started.
         */
        STARTED
    }

    /**
     * A {@link Runnable} for checking whether the process has timed out and either destroying the process
     * or scheduling another check.
     */
    private class TimeoutTask implements Runnable {

        @Override
        public void run() {
            //Because the idle timeout is reset each time the handler receives a callback, we need to
            //first check whether the _current_ timeouts have elapsed
            long nextDelay = getTimeoutDelay();
            if (nextDelay > 0L) {
                //If the timeout has been reset, reschedule this task for the new delay
                timeoutFuture = timeoutExecutor.schedule(this, nextDelay, TimeUnit.MILLISECONDS);
            } else {
                //If the timeout has elapsed, terminate the process (if it's still running)
                timeoutFuture = null; //Do this first so onExit won't try to cancel it
                process.destroyIfRunning();
            }
        }
    }
}
