package com.atlassian.bitbucket.mesh.test;

import com.google.common.collect.Lists;
import org.hamcrest.Matcher;
import org.hamcrest.Matchers;
import org.mockito.ArgumentCaptor;
import org.mockito.ArgumentMatcher;
import org.mockito.internal.hamcrest.HamcrestArgumentMatcher;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

import javax.annotation.Nonnull;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

public class MockitoUtils {

    @Nonnull
    @SafeVarargs
    @SuppressWarnings({"rawtypes", "unchecked"})
    public static <T> ArgumentMatcher<T> anyOneOf(T... values) {
        //Casting directly to Matcher<T>[] fails with a weird error, so we use the raw type and then
        //use a separate unchecked cast on the returned HamcrestArgumentMatcher below
        Matcher[] matchers = Stream.of(values)
                .map(Matchers::equalTo)
                .toArray(Matcher[]::new);

        return new HamcrestArgumentMatcher<>(Matchers.anyOf(matchers));
    }

    public static <T> T getCaptorOfType(Class<T> type, ArgumentCaptor<?> captor) {
        for (Object value : captor.getAllValues()) {
            if (value.getClass().isAssignableFrom(type)) {
                return type.cast(value);
            }
        }
        return null;
    }

    public static <T> List<T> getCaptorsOfType(Class<T> type, ArgumentCaptor<?> captor) {
        List<T> values = Lists.newArrayList();
        for (Object value : captor.getAllValues()) {
            if (value.getClass().isAssignableFrom(type)) {
                values.add(type.cast(value));
            }
        }
        return values;
    }

    //This is a wrapper around Hamcrest's containsInAnyOrder, which returns an Matcher<Iterable> and can't
    //be used if a method takes a more specific type like Collection, Set or List. This makes the return
    //type generic so the compiler can infer a more specific Iterable subclass
    @Nonnull
    @SafeVarargs
    public static <T, I extends Iterable<T>> ArgumentMatcher<I> iterableOf(T... values) {
        return new ArgumentMatcher<I>() {

            private final ArgumentMatcher<Iterable<? extends T>> delegate =
                    new HamcrestArgumentMatcher<>(Matchers.containsInAnyOrder(values));

            @Override
            public boolean matches(I t) {
                return delegate.matches(t);
            }

            @Override
            public String toString() {
                return delegate.toString();
            }
        };
    }

    public static <K, V> Answer<Void> putsValueInMap(Map<K, V> map) {
        return putsValueInMap(map, 0, 1);
    }

    @SuppressWarnings("unchecked")
    public static <K, V> Answer<Void> putsValueInMap(Map<K, V> map, int keyIndex, int valueIndex) {
        return invocation -> {
            map.put((K) invocation.getArguments()[keyIndex], (V) invocation.getArguments()[valueIndex]);
            return null;
        };
    }

    public static <K, V> Answer<Void> removesValueInMap(Map<K, V> map) {
        return removesValueInMap(map, 0);
    }

    @SuppressWarnings("SuspiciousMethodCalls")
    public static <K, V> Answer<Void> removesValueInMap(Map<K, V> map, int keyIndex) {
        return invocation -> {
            map.remove(invocation.getArguments()[keyIndex]);
            return null;
        };
    }

    public static <T> Answer<T> returnArg(int index) {
        //noinspection unchecked
        return invocationOnMock -> (T) invocationOnMock.getArguments()[index];
    }

    /**
     * @param <T> the return type of the Answer
     * @return an Answer that returns the first argument
     */
    public static <T> Answer<T> returnFirst() {
        return returnArg(0);
    }

    public static SelfAnswer returnSelf(Class<?> mockType) {
        return new SelfAnswer(mockType);
    }

    public static SelfAnswer returnSelf(Class<?> mockType, Answer<Object> defaultAnswer) {
        return new SelfAnswer(mockType, defaultAnswer);
    }

    public static Answer<Object> returnsSelf() {
        return InvocationOnMock::getMock;
    }

    public static <K, V> Answer<V> returnsValueFromMap(Map<K, V> map) {
        return returnsValueFromMap(map, 0);
    }

    @SuppressWarnings("SuspiciousMethodCalls")
    public static <K, V> Answer<V> returnsValueFromMap(Map<K, V> map, int keyIndex) {
        return invocation -> map.get(invocation.getArguments()[keyIndex]);
    }
}
