package com.atlassian.bitbucket.internal.branch.list.aheadbehind;

import com.atlassian.bitbucket.commit.graph.CommitGraphContext;
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.repository.Ref;

import javax.annotation.Nonnull;

/**
 * A graph traversal callback which compares two refs, counting how many commits ahead and behind
 * the alternate ref is from the base ref.
 *
 * @since 2.11
 */
public class AheadBehindTraversalCallback extends TraversalCallback {

    private final CommitGraphContext baseRefContext;
    private final CommitGraphContext refContext;
    private final AheadBehindCount.Builder count;
    private final Ref ref;

    /**
     * @param baseRef the base ref to which the ref is compared
     * @param ref     the ref to calculate the ahead/behind count for
     */
    public AheadBehindTraversalCallback(Ref baseRef, Ref ref) {
        baseRefContext = new CommitGraphContext.Builder().include(baseRef.getLatestCommit()).build();
        refContext = new CommitGraphContext.Builder().include(ref.getLatestCommit()).build();
        count = new AheadBehindCount.Builder();
        this.ref = ref;
    }

    @Override
    public TraversalStatus onNode(@Nonnull CommitGraphNode node) {
        if (baseRefContext.visit(node)) {
            if (refContext.visit(node)) {
                // Common ancestor found. Exclude all of its ancestors as they will be common as well
                baseRefContext.exclude(node.getParents());
                refContext.exclude(node.getParents());
            } else {
                count.incrementBehind();
            }
        } else {
            if (refContext.visit(node)) {
                count.incrementAhead();
            }
            // else the commit is unrelated.
        }

        return baseRefContext.isTraversing() || refContext.isTraversing() ?
                TraversalStatus.CONTINUE :
                TraversalStatus.FINISH;
    }

    public AheadBehindCount getCount() {
        if (baseRefContext.isAnyNodeVisited() && refContext.isAnyNodeVisited()) {
            return count.build();
        } else {
            return null;
        }
    }

    public Ref getRef() {
        return ref;
    }
}
