package org.xbib.datastructures.immutable;

import java.util.Comparator;
import java.util.Objects;
import java.util.Spliterator;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.IntConsumer;
import java.util.function.IntFunction;
import java.util.stream.IntStream;

public class Spliterators {

    private Spliterators() {
    }

    /**
     * Returns a {@code Spliterator} over the elements of {@code fromSpliterator} mapped by {@code
     * function}.
     */
    public static <InElementT extends Object, OutElementT extends Object> Spliterator<OutElementT> map(
            Spliterator<InElementT> fromSpliterator,
            Function<? super InElementT, ? extends OutElementT> function) {
        Objects.requireNonNull(fromSpliterator);
        Objects.requireNonNull(function);
        return new Spliterator<OutElementT>() {

            @Override
            public boolean tryAdvance(Consumer<? super OutElementT> action) {
                return fromSpliterator.tryAdvance(
                        fromElement -> action.accept(function.apply(fromElement)));
            }

            @Override
            public void forEachRemaining(Consumer<? super OutElementT> action) {
                fromSpliterator.forEachRemaining(fromElement -> action.accept(function.apply(fromElement)));
            }

            @Override
            public Spliterator<OutElementT> trySplit() {
                Spliterator<InElementT> fromSplit = fromSpliterator.trySplit();
                return (fromSplit != null) ? map(fromSplit, function) : null;
            }

            @Override
            public long estimateSize() {
                return fromSpliterator.estimateSize();
            }

            @Override
            public int characteristics() {
                return fromSpliterator.characteristics()
                        & ~(Spliterator.DISTINCT | Spliterator.NONNULL | Spliterator.SORTED);
            }
        };
    }


    /**
     * Returns a {@code Spliterator} that iterates over the elements of the spliterators generated by
     * applying {@code function} to the elements of {@code fromSpliterator}.
     */
    public static <InElementT extends Object, OutElementT extends Object> Spliterator<OutElementT> flatMap(
            Spliterator<InElementT> fromSpliterator,
            Function<? super InElementT, Spliterator<OutElementT>> function,
            int topCharacteristics,
            long topSize) {
        if (!((topCharacteristics & Spliterator.SUBSIZED) == 0)) {
            throw new IllegalArgumentException("flatMap does not support SUBSIZED characteristic");
        }
        if (!((topCharacteristics & Spliterator.SORTED) == 0)) {
            throw new IllegalArgumentException("flatMap does not support SORTED characteristic");
        }
        Objects.requireNonNull(fromSpliterator);
        Objects.requireNonNull(function);
        return new FlatMapSpliteratorOfObject<>(null, fromSpliterator, function, topCharacteristics, topSize);
    }

    public static class FlatMapSpliteratorOfObject<InElementT extends Object, OutElementT extends Object>
            extends FlatMapSpliterator<InElementT, OutElementT, Spliterator<OutElementT>> {

        FlatMapSpliteratorOfObject(Spliterator<OutElementT> prefix,
                                   Spliterator<InElementT> from,
                                   Function<? super InElementT, Spliterator<OutElementT>> function,
                                   int characteristics,
                                   long estimatedSize) {
            super(prefix, from, function, FlatMapSpliteratorOfObject::new, characteristics, estimatedSize);
        }
    }

    static abstract class FlatMapSpliterator<
            InElementT extends Object,
            OutElementT extends Object,
            OutSpliteratorT extends Spliterator<OutElementT>>
            implements Spliterator<OutElementT> {
        /**
         * Factory for constructing {@link FlatMapSpliterator} instances.
         */
        @FunctionalInterface
        interface Factory<InElementT extends Object, OutSpliteratorT extends Spliterator<?>> {
            OutSpliteratorT newFlatMapSpliterator(OutSpliteratorT prefix,
                                                  Spliterator<InElementT> fromSplit,
                                                  Function<? super InElementT, OutSpliteratorT> function,
                                                  int splitCharacteristics,
                                                  long estSplitSize);
        }

        OutSpliteratorT prefix;
        final Spliterator<InElementT> from;
        final Function<? super InElementT, OutSpliteratorT> function;
        final Factory<InElementT, OutSpliteratorT> factory;
        int characteristics;
        long estimatedSize;

        FlatMapSpliterator(OutSpliteratorT prefix,
                           Spliterator<InElementT> from,
                           Function<? super InElementT, OutSpliteratorT> function,
                           Factory<InElementT, OutSpliteratorT> factory,
                           int characteristics,
                           long estimatedSize) {
            this.prefix = prefix;
            this.from = from;
            this.function = function;
            this.factory = factory;
            this.characteristics = characteristics;
            this.estimatedSize = estimatedSize;
        }
        /*
         * The tryAdvance and forEachRemaining in FlatMapSpliteratorOfPrimitive are overloads of these
         * methods, not overrides. They are annotated @Override because they implement methods from
         * Spliterator.OfPrimitive (and override default implementations from Spliterator.OfPrimitive or
         * a subtype like Spliterator.OfInt).
         */

        @Override
        public final boolean tryAdvance(Consumer<? super OutElementT> action) {
            while (true) {
                if (prefix != null && prefix.tryAdvance(action)) {
                    if (estimatedSize != Long.MAX_VALUE) {
                        estimatedSize--;
                    }
                    return true;
                } else {
                    prefix = null;
                }
                if (!from.tryAdvance(fromElement -> prefix = function.apply(fromElement))) {
                    return false;
                }
            }
        }

        @Override
        public final void forEachRemaining(Consumer<? super OutElementT> action) {
            if (prefix != null) {
                prefix.forEachRemaining(action);
                prefix = null;
            }
            from.forEachRemaining(
                    fromElement -> {
                        Spliterator<OutElementT> elements = function.apply(fromElement);
                        if (elements != null) {
                            elements.forEachRemaining(action);
                        }
                    });
            estimatedSize = 0;
        }

        @Override
        public final OutSpliteratorT trySplit() {
            Spliterator<InElementT> fromSplit = from.trySplit();
            if (fromSplit != null) {
                int splitCharacteristics = characteristics & ~Spliterator.SIZED;
                long estSplitSize = estimateSize();
                if (estSplitSize < Long.MAX_VALUE) {
                    estSplitSize /= 2;
                    this.estimatedSize -= estSplitSize;
                    this.characteristics = splitCharacteristics;
                }
                OutSpliteratorT result =
                        factory.newFlatMapSpliterator(
                                this.prefix, fromSplit, function, splitCharacteristics, estSplitSize);
                this.prefix = null;
                return result;
            } else if (prefix != null) {
                OutSpliteratorT result = prefix;
                this.prefix = null;
                return result;
            } else {
                return null;
            }
        }

        @Override
        public final long estimateSize() {
            if (prefix != null) {
                estimatedSize = Math.max(estimatedSize, prefix.estimateSize());
            }
            return Math.max(estimatedSize, 0);
        }

        @Override
        public final int characteristics() {
            return characteristics;
        }
    }

    static <T extends Object> Spliterator<T> indexed(
            int size, int extraCharacteristics, IntFunction<T> function) {
        return indexed(size, extraCharacteristics, function, null);
    }

    static <T extends Object> Spliterator<T> indexed(
            int size,
            int extraCharacteristics,
            IntFunction<T> function,
            Comparator<? super T> comparator) {
        if (comparator != null) {
            if ((extraCharacteristics & Spliterator.SORTED) == 0) {
                throw new IllegalArgumentException();
            }
        }
        class WithCharacteristics implements Spliterator<T> {
            private final Spliterator.OfInt delegate;

            WithCharacteristics(Spliterator.OfInt delegate) {
                this.delegate = delegate;
            }

            @Override
            public boolean tryAdvance(Consumer<? super T> action) {
                return delegate.tryAdvance((IntConsumer) i -> action.accept(function.apply(i)));
            }

            @Override
            public void forEachRemaining(Consumer<? super T> action) {
                delegate.forEachRemaining((IntConsumer) i -> action.accept(function.apply(i)));
            }

            @Override
            public Spliterator<T> trySplit() {
                Spliterator.OfInt split = delegate.trySplit();
                return (split == null) ? null : new WithCharacteristics(split);
            }

            @Override
            public long estimateSize() {
                return delegate.estimateSize();
            }

            @Override
            public int characteristics() {
                return Spliterator.ORDERED
                        | Spliterator.SIZED
                        | Spliterator.SUBSIZED
                        | extraCharacteristics;
            }

            @Override
            public Comparator<? super T> getComparator() {
                if (hasCharacteristics(Spliterator.SORTED)) {
                    return comparator;
                } else {
                    throw new IllegalStateException();
                }
            }
        }
        return new WithCharacteristics(IntStream.range(0, size).spliterator());
    }
}
