package com.atlassian.bitbucket.scm.cache.ssh;

import com.atlassian.bitbucket.scm.cache.git.GitUtils;
import com.atlassian.bitbucket.scm.cache.git.PackRequest;
import com.atlassian.bitbucket.scm.git.protocol.GitPackets;
import com.atlassian.bitbucket.scm.git.protocol.ssh.GitSshScmRequestHandler;
import com.atlassian.bitbucket.scm.ssh.ExitCodeCallback;
import com.atlassian.bitbucket.scm.ssh.SshScmRequest;
import com.atlassian.bitbucket.throttle.ThrottleService;
import com.atlassian.bitbucket.throttle.Ticket;
import com.atlassian.bitbucket.util.CancelState;
import com.atlassian.bitbucket.util.Timer;
import com.atlassian.bitbucket.util.TimerUtils;
import com.atlassian.util.contentcache.internal.util.PipedStreams;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Charsets;
import com.google.common.base.Throwables;
import com.google.common.io.Closer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.concurrent.NotThreadSafe;
import java.io.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.atlassian.bitbucket.scm.git.protocol.GitPackets.isDone;
import static com.atlassian.bitbucket.scm.git.protocol.GitPackets.isFlush;
import static com.google.common.base.Preconditions.checkState;

@NotThreadSafe
public class UploadPackProxy implements Closeable {

    private static final int BUFFER_SIZE = 32 * 1024;
    private static final byte[] FLUSH = GitPackets.getFlush();

    /**
     * Time in millis to sleep if the remote client doesn't have data to pump to the process.
     */
    private static final long IO_PUMP_SLEEP = 5L;
    private static final String SCM_HOSTING = "scm-hosting";

    private static final Logger log = LoggerFactory.getLogger(UploadPackProxy.class);

    private final CancelState cancelState;
    private final String command;
    private final ExecutorService executor;
    private final ThrottleService throttleService;
    private final OutputStream processInput;
    private final InputStream processOutput;
    private final PipedStreams inputPipe;
    private final PipedStreams outputPipe;
    private final ByteArrayOutputStream processError;
    private final AtomicBoolean running;
    private final GitSshScmRequestHandler requestHandler;

    private volatile Throwable exception;
    private volatile int exitCode;
    private volatile Ticket hostingTicket;
    private volatile Future<?> requestFuture;
    private volatile SshScmRequest scmRequest;
    private volatile CountDownLatch scmRequestCompleteLatch;
    private volatile UploadPackState serverState;

    public UploadPackProxy(ExecutorService executor, String command, CancelState cancelState,
                           GitSshScmRequestHandler requestHandler, ThrottleService throttleService) {
        this.cancelState = cancelState;
        this.command = command;
        this.executor = executor;
        this.throttleService = throttleService;
        this.requestHandler = requestHandler;

        inputPipe = new PipedStreams(BUFFER_SIZE);
        outputPipe = new PipedStreams(BUFFER_SIZE);
        processError = new ByteArrayOutputStream(1024);
        processInput = inputPipe.output();
        processOutput = outputPipe.input();
        running = new AtomicBoolean(false);
        serverState = UploadPackState.NOT_STARTED;
    }

    public void cancel() {
        if (serverState != UploadPackState.PACK_SENT) {
            if (scmRequest != null) {
                scmRequest.cancel();
            }
            serverState = UploadPackState.ABORTED;
        }
        close();
    }

    public void close() {
        if (serverState == UploadPackState.TERMINATED) {
            return;
        }

        try {
            if (serverState == UploadPackState.REFS_ADVERTIZED) {
                try {
                    processInput.write(FLUSH);
                    processInput.flush();
                } catch (IOException e) {
                    // can occur if the process has already terminated on its own.
                }
            } else if (serverState == UploadPackState.SHALLOW_ADVERTIZED && scmRequest != null) {
                // no clean way to abort; the server is expecting a done token and will start stream the pack
                scmRequest.cancel();
            }
        } finally {
            if (!awaitDelegateScmRequestCompletion(2, TimeUnit.SECONDS)) {
                log.debug("Timed out waiting for the delegate process to complete.");
            }

            Closer closer = Closer.create();
            closer.register(processInput);
            closer.register(processOutput);
            closer.register(processError);
            try {
                closer.close();
            } catch (IOException e) {
                log.debug("One or more process streams could not be closed", e);
            }

            serverState = UploadPackState.TERMINATED;
        }
    }

    public String getStdErr() {
        return processError.size() > 0 ? new String(processError.toByteArray(), Charsets.UTF_8) : "";
    }

    public int getExitCode() {
        return exception == null ? exitCode : 1;
    }

    public PackRequest readPackRequest(PackRequest packRequest, InputStream inputStream) throws IOException {
        int packetsRead = 0;
        byte[] packet = null;
        try {
            PackRequest.Section section;
            while ((packet = readPacket(inputStream)) != null) {
                ++packetsRead;
                section = packRequest.packet(packet);
                if ((packRequest.isShallow() && isFlush(packet)) ||
                        packRequest.isFetch() || section == PackRequest.Section.DONE) {
                    // Stop reading the pack request when either of the following occurs:
                    // - the client sends 'done' indicating that it's done sending the pack request
                    // - the client sends 'deepen<num>' followed by a flush packet (0000).
                    // - the client sends a 'have' line indicating that it's a fetch
                    return packRequest;
                }
            }

            return packRequest;
        } finally {
            if (log.isTraceEnabled()) {
                log.trace("readPackRequest: {} git packets read. Last packet was '{}'", packetsRead,
                        packet == null ? "null" : new String(packet, Charsets.UTF_8).replace("\n", "\\n"));
            }
        }
    }

    public void releaseTicket() {
        if (hostingTicket != null) {
            hostingTicket.close();
            hostingTicket = null;
        }
    }

    public void streamPack(PackRequest packRequest, UploadPackState clientState,
                           InputStream inputStream, OutputStream outputStream) throws IOException {
        if (serverState == UploadPackState.NOT_STARTED) {
            // the delegate request hasn't been started yet. Start it now and bring it in sync with packRequest
            // this involves discarding output from git-upload-pack; the discarded output has already been sent
            // from a cache.
            streamRefs(null);
            checkState(serverState == UploadPackState.REFS_ADVERTIZED, "Must be in REFS_ADVERTIZED state, was %s",
                    serverState.name());
        }
        if (serverState == UploadPackState.REFS_ADVERTIZED && packRequest.isShallow()) {
            // feed the wants and deepen to git-upload-pack; discard shallow/unshallow instructions if the client has
            // already received them from a cache.
            streamShallow(inputStream, clientState == UploadPackState.REFS_ADVERTIZED ? outputStream : null);
            checkState(serverState == UploadPackState.SHALLOW_ADVERTIZED, "Must be in SHALLOW_ADVERTIZED state, was %s",
                    serverState.name());
        }

        // at this point all pack-file negotiation is complete, but the server hasn't received all input from the client.
        // The server may not have received any 'want' lines (in the case of a shallow request it will have)
        // and will not have received any 'have' line, nor the 'done' packet from the client.
        long bytesReadFromClient = 0;
        long bytesWrittenToClient = 0;
        boolean clientDoneSentToProcess = false;
        while (!requestFuture.isDone() || processOutput.available() > 0) {
            // stop when the process has finished and processOutput has been depleted

            // try to pump the inputstream to the process until the client has sent a 'done' packet
            if (!clientDoneSentToProcess && inputStream.available() > 0) {
                byte[] packet = GitUtils.readPacket(inputStream);
                if (packet != null) {
                    clientDoneSentToProcess = isDone(packet);
                    bytesReadFromClient += packet.length;
                    processInput.write(packet);
                    processInput.flush();
                }
                // give pumping data from the client to the server priority
                continue;
            }

            // after the client is done, we can simply block on reading from the process output, which is more efficient
            // than calling .available all the time.
            if (clientDoneSentToProcess || processOutput.available() > 0) {
                byte[] packet = GitUtils.readPacket(processOutput);
                if (packet == null) {
                    continue;
                }
                outputStream.write(packet);
                outputStream.flush();
                bytesWrittenToClient += packet.length;
            } else {
                // client hasn't sent a 'done' packet yet and no output is available from the git process.
                try {
                    Thread.sleep(IO_PUMP_SLEEP);
                    if (cancelState.isCanceled() && processOutput.available() == 0) {
                        // client has disconnected before the process started returning data. Time to bail
                        throw new IOException("Remote client disconnected");
                    }
                } catch (InterruptedException e) {
                    // restore the interrupt flag
                    Thread.currentThread().interrupt();
                    break;
                }
            }
        }

        log.debug("scm-cache: stream-pack completed {} bytes read, {} bytes written", bytesReadFromClient, bytesWrittenToClient);
        setServerStateOrThrow(UploadPackState.PACK_SENT);
    }

    public void streamRefs(OutputStream outputStream) throws IOException {
        checkState(serverState == UploadPackState.NOT_STARTED, "Must be in NOT_STARTED state, was %s", serverState.name());

        startDelegateRequest();

        // git-upload-pack writes all refs + flush
        copyUntilFlush(processOutput, outputStream);
        setServerStateOrThrow(UploadPackState.REFS_ADVERTIZED);
    }

    public void streamShallow(InputStream inputStream, OutputStream outputStream) throws IOException {
        checkState(serverState == UploadPackState.REFS_ADVERTIZED || serverState == UploadPackState.NOT_STARTED);

        if (serverState == UploadPackState.NOT_STARTED) {
            streamRefs(null);
        }

        // client sends the wants, followed by shallow and/or deepen section; copy to the delegate request
        copyUntilFlush(inputStream, processInput);

        // server responds by sending list of shallow/unshallow directives, followed by flush packet
        copyUntilFlush(processOutput, outputStream);
        setServerStateOrThrow(UploadPackState.SHALLOW_ADVERTIZED);
    }

    @VisibleForTesting
    boolean awaitDelegateScmRequestCompletion(long timeout, TimeUnit timeUnit) {
        if (scmRequestCompleteLatch == null) {
            // never started, so await succeeds by default
            return true;
        }

        try {
            return scmRequestCompleteLatch.await(timeout, timeUnit);
        } catch (InterruptedException e) {
            // restore the interrupted flag
            Thread.currentThread().interrupt();

            log.debug("Interrupted while waiting for the delegate process to complete {}", e.getMessage());
            return false;
        }
    }

    private void acquireTicket() {
        if (hostingTicket == null) {
            hostingTicket = throttleService.acquireTicket(SCM_HOSTING);
        }
    }

    private void copyUntilFlush(InputStream input, OutputStream output) throws IOException {
        byte[] packet;
        while ((packet = GitUtils.readPacket(input)) != null) {
            if (output != null) {
                output.write(packet);
                output.flush();
            }
            if (isFlush(packet)) {
                return;
            }
        }
        // Something went wrong, and Git apparently closed the stream without sending a flush packet.  If the Git
        // process hasn't exited already, wait for it to do so, so that the full exit status and stderr message that it
        // returned can be gathered later (in {@link #setServerStateOrThrow()}).
        try {
            requestFuture.get();
        } catch (InterruptedException e) {
            exception = e;
        } catch (ExecutionException e) {
            exception = e.getCause();
        }
    }

    private byte[] readPacket(InputStream inputStream) throws IOException {
        byte[] packet = null;
        long start = System.currentTimeMillis();
        try {
            return packet = GitUtils.readPacket(inputStream);
        } finally {
            if (log.isTraceEnabled()) {
                long duration = System.currentTimeMillis() - start;
                if (duration > 50) {
                    log.trace("slow packet read. {}ms to read '{}'", duration,
                            packet == null ? "null" : new String(packet, Charsets.UTF_8).replace("\n", "\\n"));
                }
            }
        }
    }

    private void setServerStateOrThrow(UploadPackState state) throws IOException {
        if (exception != null) {
            Throwables.propagateIfPossible(exception, IOException.class);
            // The exception _should_ be either an IOException or a RuntimeException / Error, which will be rethrown by
            // the line above. In case the code changes and this assumption is no longer valid, rethrow the exception as
            // a RuntimeException.
            Throwables.propagate(exception);
        }

        if (exitCode > 0 && requestFuture.isDone()) {
            // throw an IOException to invalidate the cache entry that is being created (if any). The exception will be
            // rethrown on all threads that are accessing the cache entry and as a result the correct exit code and
            // error message will be sent to all clients accessing the entry
            throw new ProcessFailedException(exitCode, getStdErr());
        }

        // no errors, it's safe to update the state
        serverState = state;
    }

    private void startDelegateRequest() throws IOException {
        if (running.compareAndSet(false, true)) {
            InputStream delegateInput = inputPipe.input();
            OutputStream delegateOutput = outputPipe.output();
            ExitCodeCallback delegateExitCodeCallback = exitCode -> UploadPackProxy.this.exitCode = exitCode;

            scmRequest = requestHandler.create(command, delegateInput, delegateOutput, processError, delegateExitCodeCallback)
                    .orElseThrow(() -> new IllegalStateException("The delegate SCM request is invalid"));

            /**
             * Acquire an {@code SCM_HOSTING} ticket here, <i>before</i> spawning a new thread. This means that
             * during spikes in SSH requests, we do not create so many threads on the executor that immediately
             * block on waiting for a ticket, but instead we block here and only create a thread when the ticket it
             * needs has already been successfully acquired.
             *
             * Note that the acquired ticket MUST be released from the same thread it was acquired on, which is the
             * thread that CachingSshScmRequest.handleRequest is running on.
             */
            acquireTicket();

            try {
                scmRequestCompleteLatch = new CountDownLatch(1);
                CountDownLatch latch = new CountDownLatch(1);
                try {
                    requestFuture = executor.submit(new LatchCountDownTask(() -> {
                        try (Timer ignored = TimerUtils.start("scm-cache: delegate request " + command)) {
                            latch.countDown();
                            scmRequest.handleRequest();
                        } catch (Exception e) {
                            log.debug("Exception encountered while delegating upload-pack request", e);
                            exception = e;
                        } finally {
                            // Unblock the coupled processInput and processOutput streams, which in turn allows
                            // CachingSshUploadPackRequest.handleRequest to terminate.
                            Closer closer = Closer.create();
                            closer.register(delegateInput);
                            closer.register(delegateOutput);
                            try {
                                closer.close();
                            } catch (IOException e) {
                                log.debug("Delegate streams could not be closed");
                            }
                        }
                        return null;
                    }));
                } catch (RejectedExecutionException e) {
                    scmRequestCompleteLatch.countDown();
                    throw new IllegalStateException("Could not start delegate SCM request; too many pending requests", e);
                }

                try {
                    // wait for the delegate request to start. If the thread doesn't start for some reason (not enough
                    // threads in the threadpool), there's no point in continuing at all. We'd only block waiting on IO that
                    // will never come.
                    if (!latch.await(1, TimeUnit.MINUTES)) {
                        requestFuture.cancel(true);
                        throw new IllegalStateException("Timed out waiting for delegate SCM request to start");
                    }
                } catch (InterruptedException e) {
                    requestFuture.cancel(true);
                    // restore the interrupted flag on the current thread
                    Thread.currentThread().interrupt();
                    throw new IllegalStateException("Interrupted while waiting for delegate SCM request", e);
                }
            } catch (RuntimeException e) {
                // at this point, the delegate request if at all started should have been stopped
                releaseTicket();
                scmRequestCompleteLatch.countDown();
                throw e;
            }
        }
    }

    /**
     * Utility class that counts down the scmRequestCompleteLatch upon task completion whether normally or via
     * cancellation.
     */
    private class LatchCountDownTask extends FutureTask<Void> {

        public LatchCountDownTask(Callable<Void> callable) {
            super(callable);
        }

        @Override
        protected void done() {
            scmRequestCompleteLatch.countDown();
        }
    }

}
