package com.atlassian.bitbucket.mesh.util;

import javax.annotation.Nonnull;
import java.util.*;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Collectors;

//Note: Collections.unmodifiable* methods are used instead of Guava's Immutable* types because Guava's
//      types have a strong opinion about nulls. Since these are utility methods they shouldn't enforce
//      such limitations on their parameters (especially without documenting the restriction)

/**
 * Additional utility methods missing from {@link Collectors}. The implementation of these collectors, where
 * possible, attempts to emulate the behavior of the standard {@link Collectors} factories.
 */
public class MoreCollectors {

    private MoreCollectors() {
        throw new UnsupportedOperationException(getClass().getName() +
                " is a utility class and should not be instantiated");
    }

    /**
     * Returns a {@code Collector} that accumulates elements into an immutable list.
     * <p>
     * The {@code Stream} being collected may contain {@code null} elements and, when it does, the collected
     * list will also contain {@code null} elements.
     *
     * @param <T> the type of the input elements
     * @return a {@code Collector} which collects elements into an immutable list
     * @see Collectors#toList
     */
    @Nonnull
    public static <T> Collector<T, ?, List<T>> toImmutableList() {
        return Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList);
    }

    /**
     * Returns a {@code Collector} that accumulates elements into an immutable map whose keys and values
     * are the result of applying the provided mapping functions to the input elements.
     * <p>
     * The {@code Stream} being collected may contain {@code null} elements, if the mapping functions handle
     * {@code nulls}. However, like {@code Collectors.toMap(Function, Function)}, this collector uses a
     * {@code HashMap} to collect keys and values. Since its implementation is uses {@link Map#merge}, the
     * {@code valueMapper} <i>may not</i> produce {@code null} values, even for {@code null} inputs, or a
     * {@code NullPointerException} will be thrown. The {@code keyMapper}, on the other hand, <i>may</i>
     * return {@code null} keys, but may not return <i>duplicate</i> keys or an {@code IllegalStateException}
     * will be thrown.
     *
     * @param <T>         the type of the input elements
     * @param <K>         the output type of the key mapping function
     * @param <U>         the output type of the value mapping function
     * @param keyMapper   a mapping function to produce keys, <i>which may return {@code null}</i>
     * @param valueMapper a mapping function to produce values, <i>which may not return {@code null}</i>
     * @return a {@code Collector} which collects elements into an immutable map
     * @see Collectors#toMap(Function, Function)
     */
    @Nonnull
    public static <T, K, U> Collector<T, ?, Map<K, U>> toImmutableMap(Function<? super T, ? extends K> keyMapper,
                                                                      Function<? super T, ? extends U> valueMapper) {
        return Collectors.collectingAndThen(Collectors.toMap(keyMapper, valueMapper), Collections::unmodifiableMap);
    }

    /**
     * Returns a {@code Collector} that accumulates elements into an immutable set.
     * <p>
     * The {@code Stream} being collected may contain {@code null} elements and, when it does, the collected
     * set will also contain a {@code null} element. Unlike {@code Collectors.toSet()}, this collector attempts
     * to retain the traversal order of the stream's elements.
     *
     * @param <T> the type of the input elements
     * @return a {@code Collector} which collects elements into an immutable set
     * @see Collectors#toSet
     */
    @Nonnull
    public static <T> Collector<T, ?, Set<T>> toImmutableSet() {
        return Collectors.collectingAndThen(
                //Uses a special Collector, because Collectors.toSet() does not preserve order
                Collector.of((Supplier<Set<T>>) LinkedHashSet::new, Set::add, (left, right) -> {
                    left.addAll(right);

                    return left;
                }, Collector.Characteristics.IDENTITY_FINISH),
                Collections::unmodifiableSet);
    }
}
