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

import com.atlassian.bitbucket.experimental.event.repository.AnalyticsRepositoryCloneEvent;
import com.atlassian.bitbucket.experimental.event.repository.AnalyticsRepositoryPullEvent;
import com.atlassian.bitbucket.i18n.KeyedMessage;
import com.atlassian.bitbucket.repository.Repository;
import com.atlassian.bitbucket.request.RequestContext;
import com.atlassian.bitbucket.request.RequestManager;
import com.atlassian.bitbucket.scm.cache.ScmCacheConfig;
import com.atlassian.bitbucket.scm.cache.ScmRequestType;
import com.atlassian.bitbucket.scm.cache.git.PackRequest;
import com.atlassian.bitbucket.scm.cache.internal.CacheUtils;
import com.atlassian.bitbucket.scm.cache.util.ResettableInputStream;
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.ResourceBusyException;
import com.atlassian.bitbucket.throttle.ThrottleService;
import com.atlassian.bitbucket.util.CancelState;
import com.atlassian.bitbucket.util.Timer;
import com.atlassian.bitbucket.util.TimerUtils;
import com.atlassian.event.api.EventPublisher;
import com.atlassian.util.contentcache.*;
import com.google.common.base.Throwables;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.concurrent.ExecutorService;

public class CachingSshUploadPackRequest implements SshScmRequest, CancelState {

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

    private final ContentCacheManager cacheManager;
    private final ScmCacheConfig config;
    private final OutputStream errorStream;
    private final EventPublisher eventPublisher;
    private final ExitCodeCallback exitCodeCallback;
    private final ExecutorService executor;
    private final ResettableInputStream inputStream;
    private final OutputStream outputStream;
    private final GitSshScmRequestHandler requestHandler;
    private final RequestManager requestManager;
    private final String sshCommand;
    private final DateFormat timeFormat;
    private final ThrottleService throttleService;

    private UploadPackState clientState;

    private volatile boolean canceled;
    private volatile SshScmRequest delegate;
    private volatile UploadPackProxy uploadPackProxy;
    private volatile boolean writingToCache;

    public CachingSshUploadPackRequest(String sshCommand, InputStream inputStream, OutputStream outputStream,
                                       OutputStream errorStream, ExitCodeCallback exitCodeCallback,
                                       ContentCacheManager cacheManager, ScmCacheConfig config,
                                       EventPublisher eventPublisher, GitSshScmRequestHandler requestHandler,
                                       RequestManager requestManager, ExecutorService executor,
                                       ThrottleService throttleService) {

        this.cacheManager = cacheManager;
        this.config = config;
        this.eventPublisher = eventPublisher;
        this.throttleService = throttleService;
        this.clientState = UploadPackState.NOT_STARTED;
        this.requestHandler = requestHandler;
        this.errorStream = errorStream;
        this.executor = executor;
        this.exitCodeCallback = exitCodeCallback;
        this.inputStream = new ResettableInputStream(inputStream);
        this.outputStream = outputStream;
        this.requestManager = requestManager;
        this.sshCommand = sshCommand;
        this.timeFormat = new SimpleDateFormat("HH:mm:ss.SSS");
    }

    @Override
    public void cancel() {
        if (canceled) {
            return;
        }

        synchronized (this) {
            // double checked locking
            if (canceled) {
                return;
            }
            canceled = true;
        }

        // try to let any in-progress cache write finish. There may be readers for this cache and by letting the cache
        // write finish, those readers can complete without getting an error.
        if (writingToCache) {
            log.debug("Delegate request is still writing to the cache. Not canceling it.");
            return;
        }

        try {
            // shut down the upload-pack process if it's running (and not hitting the cache)
            if (uploadPackProxy != null) {
                Repository repository = getRepository();
                //noinspection ConstantConditions
                log.debug("{}/{} Canceling SSH git-upload-pack", repository.getProject().getKey(), repository.getSlug());
                uploadPackProxy.cancel();
            }
        } catch (Exception e) {
            log.debug("Problem canceling git-upload-pack", e);
        }
    }

    @Override
    public void cancel(@Nonnull KeyedMessage message) {
        log.info("git-upload-pack canceled: '{}'", message.getLocalisedMessage());

        cancel();
    }

    @Override
    public boolean isCanceled() {
        return canceled;
    }

    @Nonnull
    @Override
    public Repository getRepository() {
        SshScmRequest delegate = getDelegate();
        if (delegate == null) {
            throw new IllegalStateException("Delegate SshScmRequest is not defined");
        }
        return delegate.getRepository();
    }

    @Override
    public void handleRequest() throws IOException {
        log.debug("scm-cache: start upload-pack-request");
        if (canceled) {
            return;
        }

        String repoIdentifier = "[" + getRepository() + "]";

        // Server sends refs to the client. This passes through the cache and may or may not create a cache entry.
        CacheResult result = CacheResult.BYPASS;
        PackRequest packRequest = null;

        int exitCode = 0;
        String stdError = null;
        try (Timer ignored = TimerUtils.start(now() + " scm-cache: upload-pack-request")) {
            if (config.isEnabled(ScmRequestType.REFS)) {
                result = streamCachedRefs();
            } else {
                log.debug("{} scm-cache: bypassing cache for ref advertisement", repoIdentifier);
                //noinspection UnusedDeclaration
                try (Timer ignored2 = TimerUtils.start("(bypass cache) refs")) {
                    getUploadPackProxy().streamRefs(outputStream);
                    clientState = UploadPackState.REFS_ADVERTIZED;
                }
            }

            if (canceled) {
                return;
            }

            // Pack negotiation: client sends list of wants, shallow and haves. Buffer everything the client sends us;
            // in case of a cache miss, it will have to be replayed back to git-upload-pack.
            inputStream.mark();
            log.debug("{} scm-cache: refs advertized, reading pack request", repoIdentifier);
            packRequest = readPackRequest(new PackRequest());

            if (canceled) {
                return;
            }

            if (packRequest.isEmpty() && packRequest.isComplete()) {
                log.debug("{} scm-cache: client is up to date", repoIdentifier);
                return;
            }

            if (packRequest.isFetch() || !config.isEnabled(ScmRequestType.UPLOAD_PACK)) {
                // client has sent 'shallow', and/or 'haves', which indicates a fetch. Fetches are not cached, so no caching
                // needs to be provided.
                log.debug("{} scm-cache: bypassing cache for sending pack", repoIdentifier);
                result = bypassCache(packRequest, clientState);
                return;
            }

            if (packRequest.isShallow()) {
                // client has sent a 'deepen' line, which indicates that the client only wants a shallow clone
                // the server responds by sending a list of objects that will be shallow in the pack file about to be sent
                if (log.isTraceEnabled()) {
                    log.trace("{} received upload-pack request {}", repoIdentifier, packRequest.toString());
                }
                result = streamCachedShallowInfo(packRequest);

                if (canceled) {
                    return;
                }

                // after the shallow information has been sent, the client can send us a list of 'have' lines. Read the rest
                // of the pack request.
                readPackRequest(packRequest);
            }

            if (canceled) {
                return;
            }

            if (packRequest.isFetch()) {
                result = bypassCache(packRequest, clientState);
            } else {
                // no 'shallow' or 'have' received. This is a (shallow) clone, which should be passed through the cache.
                if (!packRequest.isEmpty()) {
                    if (log.isTraceEnabled()) {
                        log.trace("{} received upload-pack request {}", repoIdentifier, packRequest.toString());
                    }
                    log.debug("{} scm-cache: streaming pack through cache", repoIdentifier);
                    result = streamCachedPack(packRequest);
                    log.debug("{} scm-cache: done streaming pack through cache ({})", repoIdentifier, result);
                }
            }

            if (uploadPackProxy != null) {
                exitCode = uploadPackProxy.getExitCode();
                stdError = uploadPackProxy.getStdErr();
            }
        } catch (ProcessFailedException e) {
            // The delegate request (or the delegate request of a concurrent cache access) failed with a non-zero exit
            // code and possible an error message.
            exitCode = e.getExitCode();
            stdError = e.getStdError();
        } catch (ResourceBusyException e) {
            throw e;
        } catch (Throwable t) {
            log.debug(repoIdentifier + " Exception encountered while processing a git-upload-pack request", t);
            if (canceled) {
                Throwables.propagateIfPossible(t);
            } else {
                Throwables.propagateIfPossible(t, IOException.class);
            }
        } finally {
            log.debug("{} scm-cache: upload-pack-request of type {} {}", repoIdentifier,
                    packRequest != null ? packRequest.getRequestType() : "null",
                    canceled ? "was canceled" : (exitCode == 0 ? "completed" : "failed"));

            if (StringUtils.isNotBlank(stdError)) {
                try {
                    errorStream.write(stdError.getBytes("UTF-8"));
                    errorStream.flush();
                    log.debug("{} exited with exit code {} and stderr '{}'", sshCommand, exitCode, stdError.trim());
                } catch (IOException e) {
                    log.debug("Could not write stderr {} to SSH client: '{}'", stdError, e.getMessage());
                }
            }
            addLabels(result, packRequest);
            publishEvent(packRequest);
            if (uploadPackProxy != null) {
                uploadPackProxy.close();
                uploadPackProxy.releaseTicket();
            }
            exitCodeCallback.onExit(exitCode);
        }
    }

    @Override
    public boolean isWrite() {
        return false;
    }

    @Override
    public void sendError(@Nonnull String summary, @Nonnull String detailMessage) throws IOException {
        getDelegate().sendError(summary, detailMessage);
    }

    private void addLabels(CacheResult result, PackRequest request) {
        try {
            RequestContext context = requestManager.getRequestContext();
            if (context != null) {
                if (result != null && result != CacheResult.BYPASS) {
                    context.addLabel(result.toString());
                }
                if (request != null) {
                    if (request.isShallow()) {
                        context.addLabel("shallow");
                    }
                    context.addLabel(request.getRequestType());
                }
            }
        } catch (Exception e) {
            // this can happen when the SCM Cache plugin is disabled/uninstalled while a request is being handled
            log.debug("Could not append cache hit/miss information to request labels", e);
        }
    }

    private CacheResult bypassCache(PackRequest packRequest, UploadPackState clientState) throws IOException {
        // release any buffer that the inputstream has got internally
        inputStream.reset();
        inputStream.clearMark();
        try (Timer ignored = TimerUtils.start("(bypass cache) pack")) {
            getUploadPackProxy().streamPack(packRequest, clientState, inputStream, outputStream);
            addLabels(null, packRequest);
        }
        return CacheResult.BYPASS;
    }

    private SshScmRequest getDelegate() {
        if (delegate == null) {
            delegate = requestHandler.create(sshCommand, inputStream, outputStream, errorStream, exitCodeCallback).orElse(null);
        }
        return delegate;
    }

    private String getCacheRegion() {
        return CacheUtils.getCacheRegionForRepository(getRepository());
    }

    private UploadPackProxy getUploadPackProxy() throws IOException {
        if (uploadPackProxy == null) {
            uploadPackProxy = new UploadPackProxy(executor, sshCommand, this, requestHandler, throttleService);
        }
        return uploadPackProxy;
    }

    private String now() {
        return timeFormat.format(new Date());
    }

    @SuppressWarnings("ConstantConditions")
    private void publishEvent(PackRequest packRequest) {
        //Published after request completion so that bytes read/written and duration are available
        RequestContext context = requestManager.getRequestContext();
        Repository repository = getRepository();
        if (context != null && repository != null && packRequest != null && !"refs".equals(packRequest.getRequestType())) {
            context.addCleanupCallback(() -> eventPublisher.publish(packRequest.isFetch() ?
                    new AnalyticsRepositoryPullEvent(this, getRepository(), context) :
                    new AnalyticsRepositoryCloneEvent(this, getRepository(), context))
            );
        }
    }

    private PackRequest readPackRequest(PackRequest packRequest) throws IOException {
        try (Timer ignored = TimerUtils.start("reading pack request")) {
            return getUploadPackProxy().readPackRequest(packRequest, inputStream);
        }
    }

    private CacheResult streamCached(String taskName, String cacheKey,
                                     ContentProvider contentProvider, boolean finalStep) throws IOException {
        writingToCache = true;
        try (Timer ignored = TimerUtils.start(now() + " (through cache) " + taskName)) {
            for (int i = 0; i < 5 && !canceled; ++i) {
                ContentCache cache = cacheManager.getCache(getCacheRegion());
                CacheAccess access = cache.access(cacheKey, outputStream, contentProvider);
                try {
                    if (finalStep && CacheResult.HIT.equals(access.getResult()) && uploadPackProxy != null) {
                        log.debug("scm-cache: '{}'. terminating delegate process (streaming from cache)", taskName);
                        // the pack was in the cache -> shut down the git-upload-pack process as quickly as possible
                        // and releases any throttling tickets that are being held.
                        uploadPackProxy.close();
                        uploadPackProxy.releaseTicket();
                    }

                    CacheResult result = access.stream();
                    outputStream.flush();
                    return result;
                } catch (CacheEntryExpiredException e) {
                    // cache expired - remove from cache and try again
                    cache.remove(cacheKey);
                    log.debug("{} failed (reason: {})", taskName, e.getMessage());
                    if (access.getResult() == CacheResult.MISS) {
                        // failed to write to the cache - client disconnected before the cache entry could be completed
                        throw new IOException(taskName + " failed - the client disconnected");
                    }
                }
            }
        } finally {
            writingToCache = false;
        }
        throw new IOException(taskName + " failed");
    }

    private CacheResult streamCachedPack(PackRequest packRequest) throws IOException {
        CacheResult result = streamCached("sending pack", packRequest.buildCacheKey(), new PackContentProvider(packRequest), true);
        clientState = UploadPackState.PACK_SENT;
        return result;
    }

    private CacheResult streamCachedRefs() throws IOException {
        CacheResult result = streamCached("sending refs", CacheUtils.REFS_KEY, new RefsContentProvider(), false);
        clientState = UploadPackState.REFS_ADVERTIZED;
        return result;
    }

    private CacheResult streamCachedShallowInfo(PackRequest packRequest) throws IOException {
        CacheResult result = streamCached("sending shallow info", packRequest.buildCacheKey(), new ShallowInfoContentProvider(), false);
        clientState = UploadPackState.SHALLOW_ADVERTIZED;
        return result;
    }

    private class PackContentProvider implements ContentProvider {

        private final PackRequest packRequest;

        private PackContentProvider(PackRequest packRequest) {
            this.packRequest = packRequest;
        }

        @Override
        public Date getExpiry() {
            return config.createExpiryDate(ScmRequestType.UPLOAD_PACK);
        }

        @Override
        public void apply(final OutputStream cacheOutputStream) throws IOException {
            TimerUtils.time("stream pack to cache", () -> {
                // reads the pack from the server and streams it to outputstream
                inputStream.reset();
                getUploadPackProxy().streamPack(packRequest, clientState, inputStream, cacheOutputStream);
                return null;
            });
        }
    }

    private class RefsContentProvider implements ContentProvider {

        @Override
        public Date getExpiry() {
            return config.createExpiryDate(ScmRequestType.REFS);
        }

        @Override
        public void apply(final OutputStream cacheOutputStream) throws IOException {
            TimerUtils.time("write refs to cache", () -> {
                // read the refs from the server and stream it to the cache
                getUploadPackProxy().streamRefs(cacheOutputStream);
                return null;
            });
        }
    }

    private class ShallowInfoContentProvider implements ContentProvider {

        @Override
        public Date getExpiry() {
            return config.createExpiryDate(ScmRequestType.UPLOAD_PACK);
        }

        @Override
        public void apply(final OutputStream cacheOutputStream) throws IOException {
            TimerUtils.time("write pack negotiation to cache", () -> {
                // read the shallow info from the server and stream it to the cache
                inputStream.reset();
                getUploadPackProxy().streamShallow(inputStream, cacheOutputStream);
                inputStream.mark();
                return null;
            });
        }
    }
}
