/*
 * @author Ranides Atterwim {@literal <ranides@gmail.com>}
 * @copyright Ranides Atterwim
 * @license WTFPL
 * @url http://ranides.net/projects/assira
 */
package net.ranides.assira.collection.iterators;

import java.util.concurrent.atomic.AtomicBoolean;
import lombok.experimental.UtilityClass;

import net.ranides.assira.functional.Functions;

import java.util.*;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/**
 * Static methods working with Spliterator interface.
 *
 * @author Ranides Atterwim {@literal <ranides@gmail.com>}
 */
@UtilityClass
public final class SpliteratorUtils {

    /**
     * Returns limited stream, which stops processing at first unmatched element.
     *
     * Preserves parallelization of all intermediate pipeline steps except "limit" operation itself.
     *
     * @param stream stream
     * @param predicate predicate
     * @param <T> T
     * @return stream
     */
    public static <T> Stream<T> limit(Stream<T> stream, Predicate<? super T> predicate) {
        Spliterator<T> spliterator = stream.spliterator();
        return StreamSupport.stream(limit(spliterator, predicate), stream.isParallel()).onClose(stream::close);
    }

    /**
     * Returns limited spliterator, which stops processing at first unmatched element.
     *
     * @param spliterator spliterator
     * @param predicate predicate
     * @param <T> T
     * @return spliterator
     */
    public static <T> Spliterator<T> limit(Spliterator<T> spliterator, Predicate<? super T> predicate) {
        return new LimitSpliterator<>(spliterator, predicate);
    }

    /**
     * Creates spliterator with single provided value
     *
     * @param value value
     * @param <T> T
     * @return iterator
     */
    public static <T> Spliterator<T> singletonSpliterator(T value) {
        return new SingleSpliterator<>(value);
    }

    private static class LimitSpliterator<T> extends Spliterators.AbstractSpliterator<T> {
        private final Spliterator<T> spliterator;
        private final Predicate<? super T> condition;
        boolean active;

        public LimitSpliterator(Spliterator<T> spliterator, Predicate<? super T> condition) {
            super(spliterator.estimateSize(), spliterator.characteristics() & ~(Spliterator.SIZED | Spliterator.SUBSIZED));
            this.spliterator = spliterator;
            this.condition = condition;
            this.active = true;
        }

        @Override
        public boolean tryAdvance(Consumer<? super T> action) {
            return active
                && spliterator.tryAdvance(t -> { if ((active = condition.test(t))) { action.accept(t); }})
                && active;
        }

        @Override
        public Comparator<? super T> getComparator() {
            return spliterator.getComparator();
        }
    }

    /**
     * Returns new spliterator which maps every element using provided function.
     * It preserves "split" strategy.
     *
     * @param iterator iterator
     * @param function function
     * @param <T> T
     * @param <S> S
     * @return spliterator
     */
    public static <S,T> Spliterator<T> map(Spliterator<? extends S> iterator, Functions.EachFunction<? super S,? extends T> function) {
        return new EachAdapter<>(iterator, function);
    }

    /**
     * Returns new spliterator which maps every element using provided function.
     * It preserves "split" strategy.
     *
     * @param iterator iterator
     * @param function function
     * @param <T> T
     * @param <S> S
     * @return spliterator
     */
    public static <S,T> Spliterator<T> map(Spliterator<? extends S> iterator, Function<? super S,? extends T> function) {
        return new Adapter<>(iterator, function);
    }

    private static final class Adapter<S, T> implements Spliterator<T> {

        protected final Spliterator<S> delegate;

        protected final Function<? super S,? extends T> function;

        protected <E> Adapter(Spliterator<S> delegate, Function<? super S,? extends T> function) {
            this.delegate = delegate;
            this.function = function;
        }

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

        @Override
        public Spliterator<T> trySplit() {
            Spliterator<S> out = delegate.trySplit();
            return out == null ? null : new Adapter<>(out, function);
        }

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

        @Override
        public int characteristics() {
            return delegate.characteristics();
        }

    }

    private static final class EachAdapter<S, T> implements Spliterator<T> {

        protected final Spliterator<S> delegate;

        protected final Functions.EachFunction<? super S,? extends T> function;

        protected int index;

        protected <E> EachAdapter(Spliterator<S> delegate, Functions.EachFunction<? super S,? extends T> function) {
            this.delegate = delegate;
            this.function = function;
            this.index = 0;
        }

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

        @Override
        public Spliterator<T> trySplit() {
            Spliterator<S> out = delegate.trySplit();
            return out == null ? null : new EachAdapter<>(out, function);
        }

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

        @Override
        public int characteristics() {
            return delegate.characteristics();
        }

    }

    private static class SingleSpliterator<T> implements Spliterator<T> {

	    private final T value;

        long est = 1;

        public SingleSpliterator(T value) {
            this.value = value;
        }

        @Override
        public Spliterator<T> trySplit() {
            return null;
        }

        @Override
        public boolean tryAdvance(Consumer<? super T> consumer) {
            Objects.requireNonNull(consumer);
            if (est > 0) {
                est--;
                consumer.accept(value);
                return true;
            }
            return false;
        }

        @Override
        public void forEachRemaining(Consumer<? super T> consumer) {
            tryAdvance(consumer);
        }

        @Override
        public long estimateSize() {
            return est;
        }

        @Override
        public int characteristics() {
            int out = (value != null) ? Spliterator.NONNULL : 0;

            return out | Spliterator.SIZED | Spliterator.SUBSIZED | Spliterator.IMMUTABLE |
                    Spliterator.DISTINCT | Spliterator.ORDERED;
        }
    }
}
