package cc.redberry.pipe.util;

import cc.redberry.pipe.OutputPort;

import java.util.Comparator;
import java.util.TreeSet;

public final class OrderedOutputPort<T> implements OutputPort<T> {
    private final OutputPort<T> innerOutputPort;
    private final Indexer<? super T> indexer;
    private final TreeSet<T> sortedSet;
    private long nextIndex;

    public OrderedOutputPort(OutputPort<T> innerOutputPort, Indexer<? super T> indexer) {
        this(innerOutputPort, indexer, 0);
    }

    public OrderedOutputPort(OutputPort<T> innerOutputPort, Indexer<? super T> indexer, long startsFrom) {
        this.innerOutputPort = innerOutputPort;
        this.indexer = indexer;
        this.sortedSet = new TreeSet<>(new IndexComparator<>(indexer));
        this.nextIndex = startsFrom;
    }

    @Override
    public T take() {
        if (!sortedSet.isEmpty() && indexer.getIndex(sortedSet.first()) == nextIndex) {
            ++nextIndex;
            return sortedSet.pollFirst();
        }

        T obj;

        while ((obj = innerOutputPort.take()) != null) {
            if (indexer.getIndex(obj) == nextIndex) {
                ++nextIndex;
                return obj;
            }
            sortedSet.add(obj);
        }

        if (!sortedSet.isEmpty())
            throw new IllegalStateException("Some elements left in buffer.");

        return null;
    }

    private final static class IndexComparator<T> implements Comparator<T> {
        private final Indexer<? super T> indexer;

        private IndexComparator(Indexer<? super T> indexer) {
            this.indexer = indexer;
        }

        @Override
        public int compare(T o1, T o2) {
            return Long.compare(indexer.getIndex(o1), indexer.getIndex(o2));
        }
    }
}
