package com.atlassian.bitbucket.mesh.repair;

import com.atlassian.bitbucket.mesh.client.FetchObjectsParameters;
import com.atlassian.bitbucket.mesh.config.RepairConfig;
import com.atlassian.bitbucket.mesh.coordination.WriteLock;
import com.atlassian.bitbucket.mesh.coordination.WriteTokenRegistry;
import com.atlassian.bitbucket.mesh.event.EventPublisher;
import com.atlassian.bitbucket.mesh.execution.ExecutionLocal;
import com.atlassian.bitbucket.mesh.execution.ExecutionManager;
import com.atlassian.bitbucket.mesh.git.GitAgent;
import com.atlassian.bitbucket.mesh.git.GitPackFile;
import com.atlassian.bitbucket.mesh.git.Reflog;
import com.atlassian.bitbucket.mesh.git.builder.GitProcessBuilderFactory;
import com.atlassian.bitbucket.mesh.git.builder.indexpack.IndexPackStdioHandler;
import com.atlassian.bitbucket.mesh.git.exception.AbortedException;
import com.atlassian.bitbucket.mesh.git.exception.CreateRefFailedException;
import com.atlassian.bitbucket.mesh.git.exception.DeleteRefFailedException;
import com.atlassian.bitbucket.mesh.git.exception.RepositoryRepairFailedException;
import com.atlassian.bitbucket.mesh.git.gc.GarbageCollectionManager;
import com.atlassian.bitbucket.mesh.git.gc.GcPause;
import com.atlassian.bitbucket.mesh.git.porcelain.GitPorcelain;
import com.atlassian.bitbucket.mesh.git.porcelain.StreamReflogsParameters;
import com.atlassian.bitbucket.mesh.git.porcelain.StreamRefsParameters;
import com.atlassian.bitbucket.mesh.git.ref.*;
import com.atlassian.bitbucket.mesh.git.transcode.TranscodeService;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.process.handler.LineDecoder;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.transaction.TransactionManager;
import com.atlassian.bitbucket.mesh.util.ShaUtils;
import com.atlassian.util.profiling.Ticker;
import com.google.common.base.Throwables;
import com.google.common.io.Closer;
import com.google.protobuf.ByteString;
import org.apache.commons.lang3.mutable.MutableInt;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.io.Writer;
import java.nio.ByteBuffer;
import java.nio.file.*;
import java.nio.file.attribute.BasicFileAttributes;
import java.time.Duration;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import static com.atlassian.bitbucket.mesh.git.GitConstants.PATH_LOGS;
import static com.atlassian.bitbucket.mesh.git.GitConstants.PATH_OBJECTS;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static java.lang.String.format;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.nio.file.StandardOpenOption.*;
import static java.util.Objects.requireNonNull;
import static org.apache.commons.lang3.ObjectUtils.min;

public class DefaultInteractiveRepairHelper implements InteractiveRepairHelper {

    private static final ByteBuffer EMPTY = ByteBuffer.allocate(0).asReadOnlyBuffer();

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

    private final GitAgent agent;
    private final GitProcessBuilderFactory builderFactory;
    private final RepairConfig config;
    private final EventPublisher eventPublisher;
    private final ExecutorService executorService;
    private final GarbageCollectionManager garbageCollectionManager;
    private final RepairGate gate;
    private final GitPorcelain porcelain;
    private final RefManager refManager;
    private final RepairMetrics repairMetrics;
    private final RepositoryAttributeManager repositoryAttributeManager;
    private final RepositoryHashService repositoryHashService;
    private final RepositoryManager repositoryManager;
    private final ExecutionLocal<String> repositoryUnderRepair;
    private final TransactionManager transactionManager;
    private final TranscodeService transcodeService;
    private final WriteTokenRegistry writeTokenRegistry;

    public DefaultInteractiveRepairHelper(GitAgent agent, GitProcessBuilderFactory builderFactory, RepairConfig config,
                                          EventPublisher eventPublisher, ExecutionManager executionManager,
                                          ExecutorService executorService,
                                          GarbageCollectionManager garbageCollectionManager, RepairGate gate,
                                          GitPorcelain porcelain, RefManager refManager,
                                          @Nullable RepairMetrics repairMetrics,
                                          RepositoryAttributeManager repositoryAttributeManager,
                                          RepositoryHashService repositoryHashService,
                                          RepositoryManager repositoryManager,
                                          TransactionManager transactionManager,
                                          TranscodeService transcodeService, WriteTokenRegistry writeTokenRegistry) {
        this.agent = agent;
        this.builderFactory = builderFactory;
        this.config = config;
        this.eventPublisher = eventPublisher;
        this.executorService = executorService;
        this.garbageCollectionManager = garbageCollectionManager;
        this.gate = gate;
        this.porcelain = porcelain;
        this.refManager = refManager;
        this.repairMetrics = repairMetrics;
        this.repositoryAttributeManager = repositoryAttributeManager;
        this.repositoryHashService = repositoryHashService;
        this.repositoryManager = repositoryManager;
        this.transactionManager = transactionManager;
        this.transcodeService = transcodeService;
        this.writeTokenRegistry = writeTokenRegistry;

        repositoryUnderRepair = executionManager.newExecutionLocal();
    }

    @Override
    public boolean isRepairing(@Nonnull Repository repository) {
        return requireNonNull(repository, "repository").getId().equals(repositoryUnderRepair.get());
    }

    @Nonnull
    @Override
    public InteractiveRepair startRepair(@Nonnull String repositoryId, @Nonnull Duration timeout) {
        requireNonNull(repositoryId, "repositoryId");
        requireNonNull(timeout, "timeout");

        Instant deadline = Instant.now().plus(timeout);
        //A zero or negative throttle timeout means that there's no timeout
        Duration throttleTimeout = config.getThrottleTimeout().getSeconds() > 0 ? config.getThrottleTimeout() :
                Duration.ofDays(356L);

        RepairTicket ticket = gate.acquireTicket(repositoryId, min(throttleTimeout, timeout));
        return new DefaultInteractiveRepair(repositoryId, deadline, ticket);
    }

    private class DefaultInteractiveRepair implements InteractiveRepair {

        private final Instant deadline;
        private final String repositoryId;
        private final RepairTicket ticket;
        private final RepairMetrics.CallTracer tracer;

        private volatile Set<RepairCapability> capabilities;
        private volatile String contentHash;
        private volatile boolean created;
        private volatile FetchParametersBuilder fetchParametersBuilder;
        private volatile GcPause gcPause;
        private volatile String metadataHash;
        private volatile DefaultPackProcessor packProcessor;
        private volatile DefaultRefComparer refComparer;
        private volatile DefaultReflogComparer reflogComparer;
        private volatile Repository repository;
        private volatile String sourceContentHash;
        private volatile String sourceMetadataHash;
        private volatile Path stagedReflogsDir;

        DefaultInteractiveRepair(String repositoryId, Instant deadline, RepairTicket ticket) {
            this.deadline = deadline;
            this.repositoryId = repositoryId;
            this.ticket = ticket;

            tracer = repairMetrics == null ? null : repairMetrics.startTrace();
            capabilities = EnumSet.noneOf(RepairCapability.class);
            repository = repositoryManager.findById(repositoryId);
            repositoryUnderRepair.set(repositoryId);
        }

        @Override
        public void close() {
            try (Closer closer = Closer.create()) {
                closer.register(fetchParametersBuilder);
                closer.register(gcPause);
                closer.register(packProcessor);
                closer.register(refComparer);
                closer.register(reflogComparer);
                closer.register(ticket);
                repositoryUnderRepair.remove();
                deleteStagedReflogs();
            } catch (IOException | RuntimeException e) {
                log.warn("[{}] Error while closing repair resource", repositoryId, e);
            } finally {
                if (tracer != null) {
                    tracer.stop(isConsistent());
                }
            }
        }

        @Nonnull
        @Override
        public ReflogComparer compareReflogs() {
            checkState(reflogComparer == null, "Reflog comparison should only be started once");

            Ticker ticker = tracer == null ? Ticker.NO_OP : tracer.compareReflogs();
            reflogComparer = new DefaultReflogComparer(repository, false, ticker, deadline);

            return reflogComparer;
        }

        @Nonnull
        @Override
        public RefComparer compareRefs() {
            checkState(refComparer == null, "Ref comparison should only be started once");

            fetchParametersBuilder = new FetchParametersBuilder(builderFactory, repository, created, 500);
            fetchParametersBuilder.setTimeout(getTimeout());
            Ticker ticker = tracer == null ? Ticker.NO_OP : tracer.compareRefs();
            refComparer = new DefaultRefComparer(fetchParametersBuilder, repository, ticker, deadline);

            return refComparer;
        }

        @Nonnull
        @Override
        public Set<RepairCapability> getCapabilities() {
            return capabilities;
        }

        @Override
        public FetchObjectsParameters getFetchParameters() {
            checkState(refComparer != null, "Ref comparison has not been started yet");

            return fetchParametersBuilder.maybeBuild();
        }

        @Nonnull
        @Override
        public Set<String> getNextReflogsToRepairBatch(int maxSize) {
            Set<String> result = new HashSet<>(maxSize, 1.0F);

            while (result.size() < maxSize) {
                Reflog reflog = reflogComparer.nextReflogToRepair();
                if (reflog == null) {
                    break;
                }

                if (reflog.getSize() == 0L) {
                    //The reflog file is missing, but is empty. Rather than requesting the reflog, which
                    //will not send us any data we don't already have, just create it locally
                    try {
                        startReflogRepair(reflog.getRefId())
                                .close();
                    } catch (IOException e) {
                        log.warn("[{}] Failed to create empty reflog for {}", repositoryId, reflog.getRefId());
                    }
                } else {
                    //Otherwise, if the reflog isn't empty, we'll need to download it
                    result.add(reflog.getRefId());
                }
            }

            return result;
        }

        @Override
        public Repository getRepository() {
            return repository;
        }

        @Override
        public long incrementAndGetVersion() {
            return repositoryAttributeManager.incrementAndGetVersion(repositoryId);
        }

        @Override
        public boolean isConsistent() {
            //The repository is consistent if its hashes matches the source's hashes
            boolean consistent = sourceContentHash != null && sourceContentHash.equals(contentHash) &&
                    sourceMetadataHash != null && sourceMetadataHash.equals(metadataHash);
            if (!consistent && sourceContentHash != null && sourceMetadataHash != null) {
                log.debug("[{}] Repository is not consistent.\n\tMetadata: {} vs {}\n\tContent: {} vs {}",
                        repositoryId, sourceMetadataHash, metadataHash, sourceContentHash, contentHash);
            }

            return consistent;
        }

        @Override
        public boolean isContentOutdated() {
            if (repository == null) {
                return true;
            }
            if (contentHash == null) {
                contentHash = repositoryHashService.getContentHash(repository);
            }
            return !Objects.equals(contentHash, sourceContentHash);
        }

        @Override
        public boolean isMetadataOutdated() {
            if (repository == null) {
                return true;
            }
            if (metadataHash == null) {
                metadataHash = repositoryHashService.getMetadataHash(repository);
            }
            return !Objects.equals(metadataHash, sourceMetadataHash);
        }

        @Override
        public boolean isOutOfSyncRefDetected() {
            return refComparer != null && refComparer.outOfSync;
        }

        @Override
        public void lockGc() {
            gcPause = garbageCollectionManager.pause(repository, getTimeout());
            if (gcPause == null) {
                throw new RepositoryRepairFailedException("Failed to acquire GC lock");
            }
        }

        @Override
        public boolean maybeUpdateRefsAndReflogs() {
            //First publish any staged reflogs; if this goes wrong, we don't want to update the refs yet because
            //once the refs are in sync, the reflogs will no longer be compared on a next repair attempt.
            boolean reflogsPublished = maybePublishStagedReflogs();
            boolean refsUpdates = maybeUpdateRefs();
            return refsUpdates || reflogsPublished;
        }

        @Nonnull
        @Override
        public PackProcessor newPackProcessor() {
            Ticker ticker = tracer == null ? Ticker.NO_OP : tracer.fetchObjects();
            return packProcessor = new DefaultPackProcessor(repository, ticker, deadline);
        }

        @Override
        public void resetRepositoryHashes() {
            if (repository != null) {
                repositoryHashService.resetMetadataHash(repository);
                repositoryHashService.resetContentHash(repository);
            }
        }

        @Override
        public void setCapabilities(@Nonnull Set<RepairCapability> capabilities) {
            this.capabilities = requireNonNull(capabilities, "capabilities");
        }

        @Override
        public void setSourceState(@Nonnull String sourceMetadataHash, @Nonnull String sourceContentHash) {
            this.sourceMetadataHash = requireNonNull(sourceMetadataHash, "sourceMetadataHash");
            this.sourceContentHash = requireNonNull(sourceContentHash, "sourceContentHash");

            if (repository != null) {
                //Retrieve or refresh the repository hashes
                contentHash = repositoryHashService.getContentHash(repository);
                metadataHash = repositoryHashService.getMetadataHash(repository);
            }
        }

        @Nonnull
        @Override
        public ReflogRepair startReflogRepair(@Nonnull String refId) throws IOException {
            Path path = getStagedReflogsDir().resolve(refId);

            return new DefaultReflogRepair(path, refId, fetchParametersBuilder);
        }

        @Override
        public void updateMetadata(String defaultBranch, @Nonnull Map<String, String> metadata,
                                   boolean transcodeEnabled) {
            if (repository == null) {
                //The repository does not exist locally yet, so create it
                repository = createRepository(repositoryId, defaultBranch, metadata);
                created = true;
                if (transcodeEnabled) {
                    setTranscodeEnabled(true);
                }
            } else {
                //Repository exists already, but some fields did not match. Update them all
                if (defaultBranch != null) {
                    setDefaultBranch(defaultBranch);
                }
                setMetadata(metadata);
                setTranscodeEnabled(transcodeEnabled);
            }
        }

        @Override
        public void updateVersion(long version) {
            repositoryAttributeManager.setVersion(repositoryId, version);
        }

        @Override
        public void withLock(@Nonnull Runnable action) {
            requireNonNull(action, "action");

            if (repository == null) {
                //The repository does not exist locally, so we can't lock it
                action.run();
            } else {
                try (WriteLock lock = writeTokenRegistry.startLocalLock(repository, "repair/" + UUID.randomUUID())) {
                    lock.await();
                    action.run();
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw new RepositoryRepairFailedException(
                            "Interrupted while acquiring the write lock after repair");
                }
            }
        }

        private Repository createRepository(String repositoryId, String defaultBranch, Map<String, String> metadata) {
            log.debug("[{}] Creating repository that was missing locally", repositoryId);
            return transactionManager.executeLocalOnly(tx ->
                    repositoryManager.create(new CreateRepositoryParameters.Builder(repositoryId)
                            .defaultBranch(defaultBranch)
                            .metadata(metadata)
                            .timeout(getTimeout().getSeconds())
                            .build()));
        }

        private void deleteStagedReflogs() {
            if (stagedReflogsDir != null) {
                MoreFiles.deleteQuietly(stagedReflogsDir);
            }
        }

        private Path getStagedReflogsDir() {
            if (stagedReflogsDir == null) {
                try {
                    stagedReflogsDir = Files.createTempDirectory(repository.getPath(), "logs-repair-");
                } catch (IOException e) {
                    throw new RepositoryRepairFailedException(
                            "Failed to create temporary directory for staged reflogs", e.getCause());
                }
            }
            return stagedReflogsDir;
        }

        private Duration getTimeout() {
            return Duration.between(Instant.now(), deadline);
        }

        private boolean maybePublishStagedReflogs() {
            boolean updated = false;
            if (stagedReflogsDir != null) {
                int count = publishStagedReflogs(stagedReflogsDir, repository.getPath().resolve(PATH_LOGS));
                if (count > 0) {
                    log.debug("[{}] Published {} staged reflogs", repository, count);
                    updated = true;
                }
            }

            return updated;
        }

        private boolean maybeUpdateRefs() {
            checkState(refComparer != null, "Ref comparison has not been started yet");

            List<ProposedRefUpdate> refUpdates = refComparer.getRefUpdates();
            if (refUpdates.isEmpty()) {
                return false;
            }

            //Split the update into two parts to gracefully handle invalid dark refs.
            //
            //Use git update-ref to update the public refs, but use straight file system updates to update the
            //dark refs. This ensures that we recreate the dark refs on the repair target as they are on the repair
            //source, even if they point to invalid hashes. This is better than having the repair fail.
            List<ProposedRefUpdate> publicRefUpdates = new ArrayList<>();
            for (ProposedRefUpdate refUpdate : refUpdates) {
                if (refUpdate.isDark()) {
                    updateDarkRef(refUpdate);
                    continue;
                }

                publicRefUpdates.add(refUpdate);
                if (publicRefUpdates.size() >= config.getUpdateRefBatchSize()) {
                    transactionManager.executeLocalOnlyWithoutResult(tx ->
                            refManager.updateRefs(new UpdateRefsParameters.Builder(repository)
                                    .dereferenceSymbolicRefs(false)
                                    .timeout(getTimeout().getSeconds())
                                    .transactional(true)
                                    .updates(publicRefUpdates)
                                    .build()));
                    publicRefUpdates.clear();
                }
            }

            if (!publicRefUpdates.isEmpty()) {
                transactionManager.executeLocalOnlyWithoutResult(tx ->
                        refManager.updateRefs(new UpdateRefsParameters.Builder(repository)
                                .dereferenceSymbolicRefs(false)
                                .timeout(getTimeout().getSeconds())
                                .transactional(true)
                                .updates(publicRefUpdates)
                                .build()));
            }
            return true;
        }

        private int publishStagedReflogs(Path sourceDir, Path targetDir) {
            MutableInt count = new MutableInt();

            try {
                Files.walkFileTree(sourceDir, new SimpleFileVisitor<Path>() {

                    @Override
                    public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
                        if (attrs.isRegularFile()) {
                            Path target = targetDir.resolve(sourceDir.relativize(file));
                            //Ensure the parent dir exists
                            MoreFiles.mkdir(target.getParent());
                            Files.move(file, target, StandardCopyOption.REPLACE_EXISTING);
                            count.increment();
                        }
                        return FileVisitResult.CONTINUE;
                    }
                });
            } catch (IOException e) {
                throw new RepositoryRepairFailedException("Failed to publish staged reflogs from " + sourceDir +
                        " to " + targetDir, e);
            }
            return count.intValue();
        }

        private void setDefaultBranch(String value) {
            log.debug("[{}] Updating default branch to {}", repositoryId, value);
            transactionManager.executeLocalOnlyWithoutResult(
                    tx -> repositoryManager.setDefaultBranch(repository, value));
        }

        private void setMetadata(Map<String, String> values) {
            log.debug("[{}] Updating metadata to {}", repositoryId, values);
            transactionManager.executeLocalOnlyWithoutResult(tx -> repositoryManager.setMetadata(repository, values));
        }

        private void setTranscodeEnabled(boolean value) {
            log.debug("[{}] Updating transcode setting to {}", repositoryId, value);
            transactionManager.executeLocalOnlyWithoutResult(tx -> transcodeService.setEnabled(repository, value));
        }

        private void updateDarkRef(ProposedRefUpdate refUpdate) {
            try {
                if (refUpdate.getType() == RefUpdateType.DELETE) {
                    agent.deleteLooseRef(repository, refUpdate.getRefId());
                } else {
                    agent.createRef(repository, refUpdate.getRefId(), refUpdate.getNewHash());
                }
            } catch (CreateRefFailedException | DeleteRefFailedException e) {
                throw new RepositoryRepairFailedException(format("Could not update %s in %s.",
                        refUpdate.getRefId(), repository), e);
            }
        }
    }

    private class DefaultPackProcessor implements InteractiveRepair.PackProcessor {

        private final Instant deadline;
        private final Repository destination;
        private final Runnable indexPackListener;
        private final IndexPackStdioHandler indexPackStdioHandler;
        private final Repository repository;
        private final Ticker ticker;

        private volatile boolean closed;
        private volatile Future<Path> indexPackFuture;

        DefaultPackProcessor(Repository repository, Ticker ticker, Instant deadline) {
            this.deadline = deadline;
            this.repository = repository;
            this.ticker = ticker;

            if (repository instanceof ManagedRepository && !agent.hasObjects(repository)) {
                //If possible, write the pack straight into the repository's shared-objects to avoid having to move it
                //as part of vacuuming. This is only safe if the repository does not contain any objects itself.
                destination = ((ManagedRepository) repository).getSharedObjects();
                //After index-pack completes, we raise an event to ensure the objects are protected from GC
                indexPackListener = () -> eventPublisher.publish(
                        new ObjectsUploadedToSharedObjectsEvent((ManagedRepository) repository));
            } else {
                destination = repository;
                indexPackListener = null;
            }
            indexPackStdioHandler = new IndexPackStdioHandler(destination.getPath().resolve(PATH_OBJECTS));
        }

        @Override
        public void close() {
            if (!closed) {
                closed = true;
                ticker.close();
                if (indexPackFuture == null) {
                    return;
                }

                indexPackStdioHandler.closeStdin();
                try {
                    //Give index-pack enough time to complete its processing; it can need to do a fair amount of
                    //processing after closing stdin (consistency checks, creating a idx file, etc).
                    Path path = indexPackFuture.get(Duration.between(Instant.now(), deadline).toMillis(),
                            TimeUnit.MILLISECONDS);
                    if (path != null) {
                        if (indexPackListener != null) {
                            indexPackListener.run();
                        }
                        new GitPackFile(path, true).close();
                    }
                } catch (ExecutionException e) {
                    Throwables.throwIfUnchecked(e.getCause());
                    throw new RepositoryRepairFailedException("Failed to process pack with missing objects",
                            e.getCause());
                } catch (IOException e) {
                    //If we can't delete the .keep file for the new pack, log a warning but don't fail repair
                    log.warn("[{}] Failed to remove keep file for new pack", repository, e);
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw new AbortedException("Interrupted while waiting for index-pack to finish", e);
                } catch (TimeoutException e) {
                    log.warn("[{}] Timed out waiting for index-pack to finish. Aborting...", repository);
                    indexPackFuture.cancel(true);
                    throw new AbortedException("Timed out waiting for index-pack to finish");
                }
            }
        }

        @Override
        public void write(@Nonnull ByteString bytes, boolean closed) {
            if (indexPackFuture == null) {
                indexPackFuture = builderFactory.builder(destination).indexPack()
                        .fixThin(true)
                        .keep(true)
                        .stdin(true)
                        .withTimeout(Duration.between(Instant.now(), deadline))
                        .build(indexPackStdioHandler)
                        .start();
            }
            indexPackStdioHandler.writeStdin(bytes, closed);
        }
    }

    private class DefaultRefComparer implements InteractiveRepair.RefComparer {

        private final Instant deadline;
        private final RefDiffer differ;
        private final Future<Void> future;
        private final List<ProposedRefUpdate> refUpdates;
        private final Repository repository;
        private final Ticker ticker;

        private volatile boolean closed;
        private volatile boolean outOfSync;
        private volatile long sourceRefs;
        private volatile long targetRefs;

        DefaultRefComparer(FetchParametersBuilder fetchParametersBuilder,
                           Repository repository, Ticker ticker, Instant deadline) {
            this.deadline = deadline;
            this.repository = repository;
            this.ticker = ticker;

            refUpdates = new ArrayList<>();
            //Create a differ with the repair target as the 'from' side (it's the current state) and the repair
            //source as the 'to' side (it's what we want to update to).
            differ = new RefDiffer(refUpdate -> {
                refUpdates.add(refUpdate);
                outOfSync = true;
                String newHash = refUpdate.getNewHash();
                if (newHash != null && !ShaUtils.isNull(newHash)) {
                    fetchParametersBuilder.includeIfMissing(newHash);
                }
                String oldHash = refUpdate.getOldHash();
                if (oldHash != null && !ShaUtils.isNull(oldHash)) {
                    //oldHash was derived from streaming local refs, so we know for sure it exists locally
                    fetchParametersBuilder.exclude(oldHash);
                }
            });

            future = executorService.submit(() -> {
                try (Differ.Sequence<Ref> fromSequence = differ.from()) {
                    porcelain.streamRefs(new StreamRefsParameters.Builder(repository)
                            .timeout(Duration.between(Instant.now(), deadline).getSeconds())
                            .darkRefs(true)
                            .publicRefs(true)
                            .stashRefs(true)
                            //don't fail if any of the refs are invalid; missing objects will be fetched
                            .failOnInvalidRefs(false)
                            .build(), ref -> {
                        if (fetchParametersBuilder.includeIfMissing(ref.getObjectId())) {
                            log.warn("[{}] Ref '{}' is invalid; missing object {} will be fetched",
                                    repository, ref.getId(), ref.getObjectId());
                        }
                        fromSequence.onNext(ref);
                        ++targetRefs;
                        return closed; //stop asap when canceled
                    });
                    return null;
                } catch (IOException e) {
                    throw new AbortedException("Streaming refs failed", e);
                }
            });
        }

        @Override
        public void close() {
            if (!closed) {
                differ.to().close();

                //Wait for local refs to also finish streaming
                try {
                    future.get(Duration.between(Instant.now(), deadline).toMillis(), TimeUnit.MILLISECONDS);
                } catch (InterruptedException interruptedException) {
                    Thread.currentThread().interrupt();
                    throw new AbortedException("Interrupted while streaming refs", interruptedException);
                } catch (ExecutionException executionException) {
                    Throwables.throwIfUnchecked(executionException.getCause());
                    throw new AbortedException("Streaming refs failed", executionException.getCause());
                } catch (TimeoutException e) {
                    throw new AbortedException("Timed out while streaming refs", e);
                } finally {
                    //At this point we're done waiting. Set canceled to true to ensure streamRefs terminates as soon
                    //as possible if it's still running
                    closed = true;
                    ticker.close();
                    log.debug("[{}] Processed {} source refs and {} target refs", repository, sourceRefs, targetRefs);
                }
            }
        }

        @Override
        public void onSourceRef(@Nonnull Ref sourceRef) {
            differ.to().onNext(sourceRef);
            ++sourceRefs;
        }

        @Nonnull
        List<ProposedRefUpdate> getRefUpdates() {
            return Collections.unmodifiableList(refUpdates);
        }
    }

    private class DefaultReflogComparer implements InteractiveRepair.ReflogComparer {

        private final Instant deadline;
        private final ReflogDiffer differ;
        private final Future<Void> future;
        private final Path logDir;
        private final Queue<Reflog> repairQueue;
        private final Repository repository;
        private final Ticker ticker;

        private volatile boolean closed;

        DefaultReflogComparer(Repository repository, boolean shared, Ticker ticker, Instant deadline) {
            this.deadline = deadline;
            this.repository = repository;
            this.ticker = ticker;

            //Create a differ with the repair target as the 'from' side (it's the current state) and the repair
            //source as the 'to' side (it's what we want to update to).
            differ = new ReflogDiffer(this::onDifference);
            if (shared) {
                checkArgument(repository instanceof ManagedRepository,
                        "Shared reflogs can only be compared for managed repositories");
                ManagedRepository managed = (ManagedRepository) repository;
                logDir = managed.getSharedObjects().getReflogDir(managed);
            } else {
                logDir = repository.getPath().resolve(PATH_LOGS);
            }
            repairQueue = new ArrayDeque<>();

            future = executorService.submit(() -> {
                try (Differ.Sequence<Reflog> targetReflogs = differ.from()) {
                    porcelain.streamReflogs(new StreamReflogsParameters.Builder(repository)
                                    .timeout(Duration.between(Instant.now(), deadline).getSeconds())
                                    .build(),
                            reflog -> {
                                targetReflogs.onNext(reflog);
                                return closed; //abort early when canceled
                            });
                    return null;
                } catch (IOException e) {
                    throw new AbortedException("Streaming reflogs failed", e);
                }
            });
        }

        @Override
        public void close() {
            if (!closed) {
                differ.to().close();

                //Wait for the local reflog stream to also finish
                try {
                    future.get(Duration.between(Instant.now(), deadline).toMillis(), TimeUnit.MILLISECONDS);
                } catch (InterruptedException interruptedException) {
                    Thread.currentThread().interrupt();
                    throw new AbortedException("Interrupted while streaming reflogs", interruptedException);
                } catch (ExecutionException executionException) {
                    Throwables.throwIfUnchecked(executionException.getCause());
                    throw new AbortedException("Streaming reflogs failed", executionException.getCause());
                } catch (TimeoutException e) {
                    throw new AbortedException("Timed out while streaming reflogs", e);
                } finally {
                    //At this point we're done waiting. Set canceled to true to ensure streamReflogs terminates as soon
                    //as possible if it's still running
                    closed = true;
                    ticker.close();
                }
            }
        }

        @Override
        public void onSourceReflog(@Nonnull Reflog sourceReflog) {
            differ.to().onNext(sourceReflog);
        }

        @Nullable
        Reflog nextReflogToRepair() {
            return repairQueue.poll();
        }

        private void onDifference(Reflog from, Reflog to) {
            if (to == null) {
                //Reflog no longer exists on the repair source side
                try {
                    Files.delete(logDir.resolve(from.getRefId()));
                } catch (FileNotFoundException | NoSuchFileException e) {
                    //It's ok if the file doesn't exist
                } catch (IOException e) {
                    log.warn("[{}] Failed to deleted stale reflog for {}", repository, from.getRefId());
                }
            } else {
                //Reflog is outdated or missing on the target, so it needs repair
                repairQueue.add(to);
            }
        }
    }

    private static class DefaultReflogRepair implements InteractiveRepair.ReflogRepair {

        //Reflog lines start with the old and new hash, separated by a space
        private static final Pattern REFLOG_LINE = Pattern.compile("^([0-9a-z]{40,}) ([0-9a-z]{40,}) ");

        private final FetchParametersBuilder fetchParametersBuilder;
        private final LineDecoder lineDecoder;
        private final Path path;
        private final String refId;
        private final Writer writer;

        private volatile boolean closed;

        DefaultReflogRepair(Path path, String refId, FetchParametersBuilder fetchParametersBuilder) throws IOException {
            this.fetchParametersBuilder = fetchParametersBuilder;
            this.path = path;
            this.refId = refId;

            lineDecoder = new LineDecoder(UTF_8, LineDecoder.Mode.LF);

            MoreFiles.mkdir(path.getParent());
            writer = Files.newBufferedWriter(path, UTF_8, CREATE, TRUNCATE_EXISTING, WRITE);
        }

        @Override
        public void abort() {
            if (!closed) {
                closed = true;
                try {
                    writer.close();
                } catch (IOException e) {
                    log.warn("Failed to close writer to reflog under repair {}", path);
                } finally {
                    MoreFiles.deleteQuietly(path);
                }
            }
        }

        @Override
        public void close() throws IOException {
            if (!closed) {
                closed = true;
                //Flush out any partial line remaining in the lineEncoder
                lineDecoder.forEach(EMPTY, true, this::onLine);
                writer.close();
            }
        }

        @Nonnull
        @Override
        public String getRefId() {
            return refId;
        }

        @Override
        public void write(@Nonnull ByteBuffer data) {
            lineDecoder.forEach(data, false, this::onLine);
        }

        private void analyseLine(String line) {
            Matcher matcher = REFLOG_LINE.matcher(line);
            if (matcher.find()) {
                //If any reflog commits are missing, we'll want to add them as includes so they'll be fetched,
                //but we don't add those we already have as excludes to prevent buffering an excessive set for
                //repositories where pull requests are used heavily
                fetchParametersBuilder.includeIfMissing(matcher.group(1), false);
                fetchParametersBuilder.includeIfMissing(matcher.group(2), false);
            }
        }

        private boolean onLine(String line, boolean truncated) {
            analyseLine(line);
            try {
                writer.write(line);
                writer.write('\n');
            } catch (IOException e) {
                throw new UncheckedIOException(e);
            }
            return true;
        }
    }
}
