package ru.yandex.webmaster3.core.util;

import java.util.Collection;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.function.BiConsumer;
import java.util.function.BinaryOperator;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Collector;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @author avhaliullin
 */
public class W3Collectors {
    public static <T> BinaryOperator<T> throwingMerger() {
        return (u, v) -> { throw new IllegalStateException(String.format("Duplicate key %s", u)); };
    }

    public static <T> BinaryOperator<T> acceptAnyMerger() {
        return (u, v) -> u;
    }

    public static <T> BinaryOperator<T> replacingMerger() {
        return (u, v) -> v;
    }

    public static <T, K, V, M extends Map<K, V>>
    Collector<T, ?, M> toMap(Function<? super T, ? extends K> keyMapper,
                             Function<? super T, ? extends V> valueMapper,
                             Supplier<M> mapSupplier,
                             BinaryOperator<V> merger) {
        return Collectors.toMap(keyMapper, valueMapper, merger, mapSupplier);
    }

    public static <T, K extends Enum<K>, V>
    Collector<T, ?, EnumMap<K, V>> toEnumMap(Function<? super T, ? extends K> keyMapper,
                                             Function<? super T, ? extends V> valueMapper,
                                             Class<K> enumClass) {
        return toMap(keyMapper, valueMapper, () -> new EnumMap<K, V>(enumClass), throwingMerger());
    }

    public static <K extends Enum<K>, V, T extends Map.Entry<K, V>>
    Collector<T, ?, EnumMap<K, V>> toEnumMap(Class<K> enumClass) {
        return toEnumMap(enumClass, throwingMerger());
    }

    public static <K extends Enum<K>, V, T extends Map.Entry<K, V>>
    Collector<T, ?, EnumMap<K, V>> toEnumMap(Class<K> enumClass, BinaryOperator<V> merger) {
        return toMap(Map.Entry::getKey, Map.Entry::getValue, () -> new EnumMap<K, V>(enumClass), merger);
    }

    public static <K, V, T extends Map.Entry<K, V>> Collector<T, ?, Map<K, List<V>>> groupByPair() {
        return groupByPair(Collectors.toList());
    }

    public static <K, V, T extends Map.Entry<K, V>, R extends Collection<V>>
    Collector<T, ?, Map<K, R>> groupByPair(Collector<V, ?, R> downstream) {
        return Collectors.groupingBy(Map.Entry::getKey, Collectors.mapping(Map.Entry::getValue, downstream));
    }

    public static <T, K extends Comparable<? super K>, V>
    Collector<T, ?, TreeMap<K, V>> toTreeMap(Function<? super T, ? extends K> keyMapper,
                                             Function<? super T, ? extends V> valueMapper,
                                             BinaryOperator<V> merger) {
        return toMap(keyMapper, valueMapper, TreeMap::new, merger);
    }

    public static <T, K extends Comparable<? super K>, V>
    Collector<T, ?, TreeMap<K, V>> toTreeMap(Function<? super T, ? extends K> keyMapper,
                                             Function<? super T, ? extends V> valueMapper) {
        return toMap(keyMapper, valueMapper, TreeMap::new, throwingMerger());
    }

    public static <K extends Comparable<? super K>, V, T extends Map.Entry<K, V>>
    Collector<T, ?, TreeMap<K, V>> toTreeMap() {
        return toTreeMap(throwingMerger());
    }

    public static <K extends Comparable<? super K>, V, T extends Map.Entry<K, V>>
    Collector<T, ?, TreeMap<K, V>> toTreeMap(BinaryOperator<V> merger) {
        return toMap(Map.Entry::getKey, Map.Entry::getValue, TreeMap::new, merger);
    }

    public static <K, V, T extends Map.Entry<K, V>>
    Collector<T, ?, HashMap<K, V>> toHashMap() {
        return toHashMap(throwingMerger());
    }

    public static <K, V, T extends Map.Entry<K, V>>
    Collector<T, ?, HashMap<K, V>> toHashMap(BinaryOperator<V> merger) {
        return toMap(Map.Entry::getKey, Map.Entry::getValue, HashMap::new, merger);
    }

    public static <T, A, M, R> Collector<T, A, R> flatMappingCollector(Function<T, Stream<M>> map, Collector<M, A, R> downstream) {
        return new Collector<T, A, R>() {
            @Override
            public Supplier<A> supplier() {
                return downstream.supplier();
            }

            @Override
            public BiConsumer<A, T> accumulator() {
                BiConsumer<A, M> downstreamAcc = downstream.accumulator();
                return (acc, it) -> {
                    map.apply(it).forEach(m -> downstreamAcc.accept(acc, m));
                };
            }

            @Override
            public BinaryOperator<A> combiner() {
                return downstream.combiner();
            }

            @Override
            public Function<A, R> finisher() {
                return downstream.finisher();
            }

            @Override
            public Set<Characteristics> characteristics() {
                return downstream.characteristics();
            }
        };
    }

    public static <T, A, R> Collector<T, A, R> filtering(Predicate<T> predicate, Collector<T, A, R> downstream) {
        return new Collector<T, A, R>() {
            @Override
            public Supplier<A> supplier() {
                return downstream.supplier();
            }

            @Override
            public BiConsumer<A, T> accumulator() {
                BiConsumer<A, T> downstreamAcc = downstream.accumulator();
                return (acc, it) -> {
                    if (predicate.test(it)) {
                        downstreamAcc.accept(acc, it);
                    }
                };
            }

            @Override
            public BinaryOperator<A> combiner() {
                return downstream.combiner();
            }

            @Override
            public Function<A, R> finisher() {
                return downstream.finisher();
            }

            @Override
            public Set<Characteristics> characteristics() {
                return downstream.characteristics();
            }
        };
    }
}
