package com.atlassian.bitbucket.mesh.util;

import com.google.common.collect.Iterators;
import com.google.common.collect.PeekingIterator;

import javax.annotation.Nonnull;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.stream.BaseStream;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import static java.util.Arrays.stream;
import static java.util.Comparator.comparing;
import static java.util.Comparator.naturalOrder;
import static java.util.Objects.requireNonNull;
import static java.util.Spliterator.SORTED;
import static java.util.Spliterators.spliteratorUnknownSize;
import static java.util.stream.Collectors.toList;

/**
 * Additional utility methods missing from {@link com.google.common.collect.Streams}.
 */
public class MoreStreams {

    /**
     * Returns a stream of the merged contents of all given streams, containing every element of the input streams.
     * Equivalent entries will not be de-duplicated. Sorting is stable, for equal elements, the order is the same as
     * the order of the streams.
     * <p>
     * Callers must ensure that the source streams are in non-descending order as this method does not sort its input.
     * If unsure, call {@link Stream#sorted()} before passing it to this method.
     *
     * @param streams the streams to merge
     * @param <T>     the type of the stream
     * @return a sorted stream that contains all elements of the input streams
     */
    @SafeVarargs
    @Nonnull
    public static <T extends Comparable<T>> Stream<T> mergeSorted(@Nonnull Stream<T>... streams) {
        requireNonNull(streams, "streams");

        Iterator<T> iterator = new MergingIterator<>(stream(streams).map(BaseStream::iterator).collect(toList()));
        return StreamSupport.stream(spliteratorUnknownSize(iterator, SORTED), false);
    }

    /**
     * Returns a stream of the merged contents of all given streams, containing every element of the input streams.
     * Equivalent entries will not be de-duplicated. Sorting is stable, for equal elements, the order is the same as
     * the order of the streams.
     * <p>
     * Callers must ensure that the source streams are in non-descending order as this method does not sort its input.
     * If unsure, call {@link Stream#sorted()} before passing it to this method.
     *
     * @param streams the streams to merge
     * @param <T>     the type of the stream
     * @return a sorted stream that contains all elements of the input streams
     */
    @Nonnull
    public static <T> Stream<T> mergeSorted(@Nonnull Iterable<Stream<T>> streams,
                                            @Nonnull Comparator<T> comparator) {
        requireNonNull(streams, "streams");
        requireNonNull(comparator, "comparator");

        Iterable<Iterator<T>> collect = streamIterable(streams).map(BaseStream::iterator).collect(toList());
        Iterator<T> iterator = new MergingIterator<>(collect, comparator);
        return StreamSupport.stream(spliteratorUnknownSize(iterator, SORTED), false);
    }

    /**
     * Converts an {@code Iterable} to a {@code Stream}.
     *
     * @param iterable the {@code Iterable} to stream
     * @param <T>      the type of element in the {@code Iterable}
     * @return a {@code Stream} over the provided {@code iterable}
     * @throws NullPointerException if the provided {@code iterable} is {@code null}
     */
    @Nonnull
    public static <T> Stream<T> streamIterable(@Nonnull Iterable<T> iterable) {
        return StreamSupport.stream(requireNonNull(iterable, "iterable").spliterator(), false);
    }

    private static class MergingIterator<T> implements Iterator<T> {

        private final Comparator<PeekingIterator<? extends T>> comparator;
        private final List<PeekingIterator<? extends T>> iterators;

        private T next;

        public MergingIterator(Iterable<? extends Iterator<? extends T>> iterators, Comparator<? super T> comparator) {
            this.iterators = MoreStreams.streamIterable(iterators)
                    .map(Iterators::peekingIterator)
                    .filter(Iterator::hasNext)
                    .collect(toList());
            this.comparator = comparing(PeekingIterator::peek, comparator);
        }

        @SuppressWarnings("unchecked")
        public <U extends Comparable<U>> MergingIterator(Iterable<? extends Iterator<? extends U>> iterators) {
            this((Iterable<? extends Iterator<? extends T>>) iterators, (Comparator<? super T>) naturalOrder());
        }

        @Override
        public boolean hasNext() {
            if (next != null) {
                return true;
            }
            return iterators.stream()
                    .min(comparator)
                    .map(min -> {
                        next = min.next();
                        if (!min.hasNext()) {
                            iterators.remove(min);
                        }
                        return true;
                    })
                    .orElse(false);
        }

        @Override
        public T next() {
            if (!hasNext()) {
                throw new NoSuchElementException();
            }
            T tmp = next;
            next = null;
            return tmp;
        }
    }
}
