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

import com.atlassian.bitbucket.commit.CommitService;
import com.atlassian.bitbucket.commit.graph.TraversalRequest;
import com.atlassian.bitbucket.repository.Ref;
import com.atlassian.bitbucket.repository.RefMetadataContext;
import com.atlassian.bitbucket.repository.RefMetadataProvider;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;

import javax.annotation.Nonnull;
import java.util.Collections;
import java.util.Map;
import java.util.Set;

public class AheadBehindRefMetadataProvider implements RefMetadataProvider<AheadBehindCount> {

    private final CommitService commitService;

    public AheadBehindRefMetadataProvider(CommitService commitService) {
        this.commitService = commitService;
    }

    @Nonnull
    @Override
    public Map<Ref, AheadBehindCount> getMetadata(@Nonnull RefMetadataContext context) {
        Set<Ref> refs = Sets.newHashSet(context.getRefs());
        refs.remove(context.getBaseRef());

        if (refs.isEmpty()) {
            return Collections.emptyMap();
        }

        BatchingAheadBehindTraversalCallback callback = BatchingAheadBehindTraversalCallback.create(context.getBaseRef(), refs);

        commitService.traverse(
                new TraversalRequest.Builder()
                        .repository(context.getRepository())
                        .include(Iterables.transform(
                                Iterables.concat(refs, Collections.singleton(context.getBaseRef())),
                                Ref::getLatestCommit)
                        )
                        .build(),
                callback);

        return callback.get();
    }
}
