package com.atlassian.stash.internal.commit.graph.cache;

import com.atlassian.bitbucket.commit.MinimalCommit;
import com.atlassian.bitbucket.commit.graph.CommitGraphNode;
import com.atlassian.bitbucket.commit.graph.TraversalCallback;
import com.atlassian.bitbucket.commit.graph.TraversalStatus;
import com.atlassian.bitbucket.commit.graph.TraversalSummary;
import com.atlassian.bitbucket.ServerException;
import com.atlassian.bitbucket.i18n.I18nService;

import javax.annotation.Nonnull;
import java.io.Closeable;
import java.io.IOException;
import java.io.OutputStream;
import java.util.Collection;

import static com.atlassian.stash.internal.commit.graph.cache.CachedCommitGraphOutputStream.CachedCommitGraphOutputStreamIOException;
import static com.google.common.collect.Collections2.transform;

/**
 * Persists the commit graph to an {@link OutputStream} as it is traversed.
 *
 * @since 2.11
 */
public class PersistingTraversalCallback extends TraversalCallback implements Closeable {

    private final OutputStream outputStream;
    private final I18nService i18nService;

    public PersistingTraversalCallback(OutputStream outputStream, I18nService i18nService) {
        // no need to buffer the output stream, this is handled by the content cache internally
        this.outputStream = outputStream;
        this.i18nService = i18nService;
    }

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

    @Override
    public TraversalStatus onNode(@Nonnull CommitGraphNode node) {
        String commitId = node.getCommit().getId();
        Collection<String> parentCommitIds = transform(node.getParents(), MinimalCommit::getId);

        try {
            outputStream.write(CachedCommitGraphUtils.getCommitIdAsByteArray(commitId));
            outputStream.write(CachedCommitGraphUtils.getCountAsByteArray(parentCommitIds.size()));
            for (String parentCommitId : parentCommitIds) {
                outputStream.write(CachedCommitGraphUtils.getCommitIdAsByteArray(parentCommitId));
            }
        } catch (CachedCommitGraphOutputStreamIOException e) {
            // The callback is connected directly to a CachedCommitGraphOutputStream. This can happen
            // when the cache is bypassed. As a result we can terminate the traversal early
            return TraversalStatus.FINISH;
        } catch (IOException e) {
            throw new ServerException(i18nService.createKeyedMessage("bitbucket.commit.graph.traversal.write.unexpectedioexception"),
                    e);
        }
        return TraversalStatus.CONTINUE;
    }

    @Override
    public void onEnd(@Nonnull TraversalSummary summary) {
        super.onEnd(summary);
        try {
            // this is fix for a bug in ContentCache that doesn't flush small streams
            outputStream.flush();
        } catch (IOException e) {
            throw new ServerException(
                    i18nService.createKeyedMessage("bitbucket.commit.graph.traversal.flush.unexpectedioexception"), e);
        }
    }
}
