package com.atlassian.bitbucket.mesh.grpc;

import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.process.OsProcess;
import com.atlassian.bitbucket.mesh.rpc.v1.RpcStreamChunk;
import com.google.protobuf.ByteString;
import com.zaxxer.nuprocess.NuProcess;

import javax.annotation.Nonnull;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.atlassian.bitbucket.mesh.io.ByteUtils.copy;
import static java.util.Objects.requireNonNull;

/**
 * Writes an incoming stream of {@link RpcStreamChunk stdin chunks} to a process stdin when that process is ready for
 * more stdin, possibly blocking until the process is ready for more stdin data.
 */
public class StdinObserver {

    private final AtomicBoolean closeRequested;
    private final ByteBuffer pendingStdin;

    private volatile OsProcess process;
    private volatile boolean stdinClosed;

    public StdinObserver() {
        closeRequested = new AtomicBoolean();
        pendingStdin = ByteBuffer.allocate(NuProcess.BUFFER_CAPACITY);
    }

    /**
     * @return {@code true} if {@code stdin} has been closed; otherwise, {@code false} if it will be closed later
     */
    public boolean closeStdin() {
        if (stdinClosed) {
            return true;
        }

        synchronized (pendingStdin) {
            //Invariant: When acquiring the lock, pendingStdin is ready for _writing_
            if (pendingStdin.position() == 0) {
                //If there's no pendingStdin, we can just close stdin immediately
                process.closeStdin(false);
                stdinClosed = true;
            } else {
                //Otherwise, we need to signal that the close has been requested
                //so a future onStdinReady callback can close it
                closeRequested.set(true);
            }

            //As a safety precaution, notify any threads blocked waiting for pendingStdin
            //so they can detect that stdin has been closed
            pendingStdin.notifyAll();
        }

        return stdinClosed;
    }

    public void onStart(@Nonnull OsProcess process) {
        this.process = requireNonNull(process, "process");
        if (pendingStdin.hasRemaining()) {
            process.wantWrite();
        }
    }

    /**
     * Writes any {@link #pendingStdin pending data} to the provided buffer to be forwarded to the process's
     * {@code stdin}. If some pending data can't be written, {@code true} will be returned to request another
     * buffer to write more.
     * <p>
     * If {@link #closeStdin} has been called, once all pending data has been written the process's {@code stdin}
     * will be closed.
     *
     * @param buffer a buffer to fill with data to write to the process's {@code stdin}
     * @return {@code true} if there is additional {@link #pendingStdin pending data} to write; otherwise,
     *         {@code false} to wait for more data from the client
     */
    public boolean onStdinReady(@Nonnull ByteBuffer buffer) {
        boolean more;
        synchronized (pendingStdin) {
            //Invariant: When acquiring the lock, pendingStdin is ready for _writing_
            pendingStdin.flip();

            if (copy(pendingStdin, buffer)) {
                pendingStdin.clear(); //Clear the buffer to prepare for writing
                more = false; //No more stdin to provide, at least for now

                //Once the pendingStdin buffer is emptied, we can action closeRequested
                //if it's been set
                if (closeRequested.compareAndSet(true, false)) {
                    process.closeStdin(false);
                    stdinClosed = true;
                }
            } else {
                pendingStdin.compact();

                //We know we have more stdin to provide, so we need another callback.
                //Note that, even if closeRequested is set, we can't close stdin yet
                more = true;
            }

            pendingStdin.notify();
        }
        buffer.flip();

        return more;
    }

    /**
     * Writes the provided {@link ByteString} to {@link #pendingStdin} and signals {@code wantWrite()} to
     * request an {@link #onStdinReady} callback.
     * <p>
     * If the provided data is too large to fit in {@link #pendingStdin}, this method will break it down
     * into smaller chunks and block until {@link #onStdinReady} signals the previous chunk has been
     * written to the process. This prevents buffering an excessive amount of data in gRPC or NuProcess.
     *
     * @param data the data to write
     */
    public void writeStdin(@Nonnull ByteString data, boolean closed) {
        requireNonNull(data, "data");

        List<ByteBuffer> stdin = data.asReadOnlyByteBufferList();
        int index = 0;
        if (stdin.stream().anyMatch(ByteBuffer::hasRemaining)) {
            while (true) {
                synchronized (pendingStdin) {
                    //Invariant: When acquiring the lock, pendingStdin is ready for _writing_
                    boolean copied = true;
                    while (index < stdin.size()) {
                        copied = copy(stdin.get(index), pendingStdin);
                        if (copied) {
                            //The stdin[index] buffer was completely consumed, move on to the next
                            index++;
                        } else {
                            //The stdin[index] buffer did not fit, stop copying
                            break;
                        }
                    }

                    //Always signal the process, after we copy some bytes into the pending
                    //buffer, so it can start preparing the onStdinReady callback
                    process.wantWrite();

                    if (copied) {
                        break;
                    }

                    //The entire chunk didn't fit, so we need to wait for space. This reduces
                    //monitor contention between onStdinReady and writeChunk
                    try {
                        //The delay here is a backstop, to guard against any edge case where we don't get
                        //notified. It's intentionally high enough that it shouldn't normally get hit
                        pendingStdin.wait(3000L);
                    } catch (InterruptedException e) {
                        throw new AbortedException("Interrupted while trying to queue stdin", e);
                    }

                    if (stdinClosed || !process.isRunning() || process.isCanceled()) {
                        //If stdin is closed while we're waiting for space to write, or if the process is no longer running,
                        //we break out
                        break;
                    }
                }
            }
        }

        if (closed) {
            closeStdin();
        }
    }
}
