package net.ranides.assira.collection.query.derived;

import lombok.RequiredArgsConstructor;
import net.ranides.assira.collection.iterators.ForwardIterator;
import net.ranides.assira.collection.iterators.IteratorUtils;
import net.ranides.assira.collection.query.CQuery;
import net.ranides.assira.collection.query.CQueryAbstract;
import net.ranides.assira.collection.query.support.BaseEach;
import net.ranides.assira.collection.query.support.BaseIterable;
import net.ranides.assira.collection.query.support.BaseState;
import net.ranides.assira.functional.Consumers;
import net.ranides.assira.functional.Predicates;

import java.util.Collections;
import java.util.Iterator;
import java.util.Stack;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
import java.util.stream.Stream;

/**
 * Stream implementation for tree traversal.
 * It uses DFS algorithm.
 *
 * It can visit any tree, as long as it can provide 3 operations:
 *   - root node
 *   - list of children of provided node
 *   - predicate returning true, if node is terminal
 *
 * @param <T> T
 */
@RequiredArgsConstructor
public class CQTreeDFS<T> extends CQueryAbstract<T> {

    protected final Supplier<T> root;

    protected final Function<T, CQuery<T>> children;

    protected final Predicate<T> isTerm;

    protected final Supplier<? extends RuntimeException> message;

    @Override
    public Supplier<? extends RuntimeException> message() {
        return message;
    }

    @Override
    public CQuery<T> message(Supplier<? extends RuntimeException> message) {
        return new CQTreeDFS<>(root, children, isTerm, message);
    }

    @Override
    public CQuery<T> parallel() {
        return this;
    }

    @Override
    public CQuery<T> sequential() {
        return this;
    }

    @Override
    public boolean isParallel() {
        return false;
    }

    @Override
    public boolean hasFastEach() {
        return true;
    }

    @Override
    public Stream<T> stream() {
        return BaseIterable.stream(this);
    }

    @Override
    public Iterator<T> iterator() {
        return IteratorUtils.withCause(new TreeIterator(root.get()), message);
    }

    @Override
    public void forEach(Consumer<? super T> consumer) {
        forEach0(root.get(), consumer);
    }

    @Override
    public void forEach(Consumers.EachConsumer<? super T> consumer) {
        BaseState.Index index = BaseState.newIndex(this);
        forEach0(root.get(), v -> consumer.accept(index.next(), v));
    }

    @Override
    public boolean whileEach(Predicate<? super T> consumer) {
        return whileEach0(root.get(), consumer);
    }

    @Override
    public boolean whileEach(Predicates.EachPredicate<? super T> consumer) {
        BaseState.Index index = BaseState.newIndex(this);
        return whileEach0(root.get(), v -> consumer.test(index.next(), v) );
    }

    private void forEach0(T node, Consumer<? super T> consumer) {
        consumer.accept(node);
        if(!isTerm.test(node)) {
            children.apply(node).forEach(child -> forEach0(child, consumer));
        }
    }

    private boolean whileEach0(T node, Predicate<? super T> consumer) {
        if(!consumer.test(node)) {
            return false;
        }
        if(!isTerm.test(node)) {
            return children.apply(node).whileEach(child -> whileEach0(child, consumer));
        }
        return true;
    }

    @Override
    public int size() {
        return BaseEach.size(this);
    }

    private final class TreeIterator extends ForwardIterator<T> {

        private final Stack<Iterator<T>> window;

        public TreeIterator(T node) {
            this.window = new Stack<>();
            this.window.add(Collections.singletonList(node).iterator());
        }

        @Override
        protected boolean next(Consumer<? super T> action) {
            while (!window.isEmpty()) {
                Iterator<T> nodes = window.peek();
                if (nodes.hasNext()) {
                    T node = nodes.next();
                    action.accept(node);
                    if(!isTerm.test(node)) {
                        window.push(children.apply(node).iterator());
                    }
                    return true;
                } else {
                    window.pop();
                }
            }
            return false;
        }
    }

}
