package com.atlassian.bitbucket.mesh.test.util.ssh;

import com.atlassian.bitbucket.mesh.git.builder.GitProcessBuilder;
import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.process.handler.AbstractProcessStdioHandler;
import org.apache.sshd.server.Environment;
import org.apache.sshd.server.ExitCallback;
import org.apache.sshd.server.channel.ChannelDataReceiver;
import org.apache.sshd.server.channel.ChannelSession;
import org.apache.sshd.server.command.Command;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.nio.channels.WritableByteChannel;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * SSH {@link Command} implementation to that starts a {@link com.atlassian.bitbucket.mesh.git.builder.GitProcess}
 * and pumps the data between the SSH client and the process.
 */
@SuppressWarnings("NonPrivateFieldAccessedInSynchronizedContext")
class GitSshCommand extends AbstractProcessStdioHandler<Void> implements Command, ChannelDataReceiver {

    private final GitProcessBuilder builder;
    private final AtomicBoolean closeRequested;
    private final ByteBuffer pendingStdin;

    private ChannelSession channel;
    private OutputStream err;
    private WritableByteChannel errChannel;
    private ExitCallback exitCallback;
    private Future<Void> future;
    private OutputStream out;
    private WritableByteChannel outChannel;

    private volatile boolean stdinClosed;

    GitSshCommand(GitProcessBuilder builder) {
        this.builder = builder;

        closeRequested = new AtomicBoolean();
        pendingStdin = ByteBuffer.allocate(64 * 1024);
    }

    @Override
    public void close() throws IOException {
        closeStdin();
    }

    @Override
    public int data(ChannelSession channel, byte[] buf, int start, int len) {
        // client sent us some data
        ByteBuffer stdin = ByteBuffer.wrap(buf, start, len);
        if (stdin.hasRemaining()) {
            while (true) {
                synchronized (pendingStdin) {
                    //Invariant: When acquiring the lock, pendingStdin is ready for _writing_
                    boolean copied = copy(stdin, pendingStdin);

                    //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 {
                        pendingStdin.wait(60000L);
                    } catch (InterruptedException e) {
                        throw new AbortedException("Interrupted while trying to queue stdin", e);
                    }

                    if (stdinClosed) {
                        //If stdin is closed while we're waiting for space to write,
                        break;
                    }
                }
            }
        }
        return len;
    }

    @Override
    public void destroy(ChannelSession channel) {
        if (future != null && !future.isDone()) {
            future.cancel(true);
            synchronized (pendingStdin) {
                closeStdin();
                pendingStdin.notifyAll();
            }
        }
    }

    @Override
    public Void getResult() {
        return null;
    }

    @Override
    public void onExit(int exitCode) {
        if (channel != null && channel.isOpen()) {
            exitCallback.onExit(exitCode);
        }
    }

    @Override
    public void onStderr(@Nonnull ByteBuffer buffer, boolean closed) {
        try {
            errChannel.write(buffer);
            err.flush();
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    @Override
    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.notifyAll();
        }
        buffer.flip();

        return more;
    }

    @Override
    public void onStdout(@Nonnull ByteBuffer buffer, boolean closed) {
        try {
            outChannel.write(buffer);
            out.flush();
        } catch (IOException e) {
            throw new UncheckedIOException(e);
        }
    }

    @Override
    public void setErrorStream(OutputStream err) {
        this.err = err;
        errChannel = Channels.newChannel(err);
    }

    @Override
    public void setExitCallback(ExitCallback callback) {
        exitCallback = callback;
    }

    @Override
    public void setInputStream(InputStream in) {
        // ignored
    }

    @Override
    public void setOutputStream(OutputStream out) {
        this.out = out;
        outChannel = Channels.newChannel(out);
    }

    @Override
    public void start(ChannelSession channel, Environment env) {
        this.channel = channel;
        channel.setDataReceiver(this);

        env.getEnv().forEach(builder::withEnvironment);
        future = builder.build(this).start();
    }

    @Override
    protected 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;
    }

    /**
     * Copies data from the {@code src} buffer into the {@code dst}, and returns whether all data
     * was copied (as a convenience; the caller could also check {@code src.hasRemaining()}).
     *
     * @param src the source buffer
     * @param dst the destination buffer
     * @return {@code true} if the {@code src} buffer was copied in full; otherwise, {@code false}
     *         if the {@code dst} buffer could not hold all the data
     */
    private static boolean copy(ByteBuffer src, ByteBuffer dst) {
        int available = dst.remaining();
        if (src.remaining() > available) {
            //The source has more data than we can fit in the destination's available space,
            //so slice off what we can and put that, then advance the position
            ByteBuffer slice = src.slice();
            slice.limit(available);
            dst.put(slice);

            src.position(src.position() + available);

            return false;
        }

        //The full source will fit in the destination
        dst.put(src);

        return true;
    }
}
