package com.atlassian.bitbucket.mockito;

import com.atlassian.bitbucket.util.Page;
import com.atlassian.bitbucket.util.PageRequest;
import com.atlassian.bitbucket.util.PageUtils;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

import javax.annotation.Nonnull;

import java.util.function.Function;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

/**
 * {@link Answer} implementation that generates pages of elements in accordance with provided {@code PageRequest} up to
 * a specified total number of elements.
 * <p>
 * Example usage:
 * <pre>
 * {@code private static final Function<Integer, Foo> FOO_FACTORY = new Function<Integer, Foo> {
 *     @Override
 *     public Foo apply(Integer position) {
 *         return new MockFoo(position);
 *     }
 * }
 *
 * // return up to 500 Foo objects, using FOO_FACTORY to create mock Foo objects
 * when(fooService.findAllFoos(eq("foo"), any(PageRequest.class))).thenAnswer(withPagesUpTo(500, FOO_FACTORY));
 * }
 * </pre>
 * <p>
 * For best results, it should not be used with other answers/returns (for the particular method invocation being
 * stubbed).
 */
public class PagesAnswer<T> implements Answer<Page<T>> {

    private final int elementCount;
    private final Function<Integer, T> generator;

    private PagesAnswer(int elementCount, Function<Integer, T> generator) {
        this.elementCount = elementCount;
        this.generator = restrictedElementGenerator(checkNotNull(generator, "generator"));
    }

    @Nonnull
    public static <T> PagesAnswer<T> withPagesUpTo(int maxElementCount, @Nonnull Function<Integer, T> generator) {
        checkArgument(maxElementCount > 0, "maxElementCount must be >0");

        return new PagesAnswer<>(maxElementCount, generator);
    }

    @Nonnull
    @Override
    public Page<T> answer(@Nonnull InvocationOnMock invocation) throws Throwable {
        PageRequest request = MockitoUtils.findPageRequest(invocation);

        return PageUtils.createPage(request, generator);
    }

    private Function<Integer, T> restrictedElementGenerator(final Function<Integer, T> original) {
        return i -> {
            if (i >= elementCount) {
                return null;
            }

            return original.apply(i);
        };
    }
}
