package com.timgroup.eventstore.stitching;

import com.google.common.collect.Iterators;
import com.google.common.collect.PeekingIterator;
import org.joda.time.DateTime;

import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;

import static java.util.stream.Collectors.toList;

public final class EventStoreStitchingIterator implements Iterator<EventInIdentifiedStream> {
    private final Clock clock;
    private final Duration delay;
    private final List<PeekingIterator<EventInIdentifiedStream>> underlying;

    private Long cutOffTime;
    private PeekingIterator<EventInIdentifiedStream> iteratorWhoseHeadIsNext;

    public EventStoreStitchingIterator(Clock clock, Duration delay, List<Iterator<EventInIdentifiedStream>> iterators) {
        this.clock = clock;
        this.delay = delay;
        this.underlying = iterators.stream().map(Iterators::peekingIterator).collect(toList());
    }

    @Override
    public boolean hasNext() {
        return getIteratorWhoseHeadIsNext() != null;
    }

    @Override
    public EventInIdentifiedStream next() {
        Iterator<EventInIdentifiedStream> candidate = getIteratorWhoseHeadIsNext();
        if (candidate == null) {
            throw new NoSuchElementException();
        }
        iteratorWhoseHeadIsNext = null;
        return candidate.next();
    }

    private Iterator<EventInIdentifiedStream> getIteratorWhoseHeadIsNext() {
        if (iteratorWhoseHeadIsNext != null) {
            return iteratorWhoseHeadIsNext;
        }
        Iterator<PeekingIterator<EventInIdentifiedStream>> streams = underlying.iterator();
        while (streams.hasNext()) {
            PeekingIterator<EventInIdentifiedStream> eventStream = streams.next();

            Instant potentialCutoffTime = clock.instant();
            if (eventStream.hasNext()) {
                EventInIdentifiedStream candidate = eventStream.peek();
                DateTime candidateTimestamp = candidate.event.effectiveTimestamp();
                if (cutOffTime != null && candidateTimestamp.isAfter(cutOffTime)) {
                    streams.remove();
                } else if (iteratorWhoseHeadIsNext == null || candidateTimestamp.isBefore(iteratorWhoseHeadIsNext.peek().event.effectiveTimestamp())) {
                    iteratorWhoseHeadIsNext = eventStream;
                }
            } else {
                streams.remove();
                if (this.cutOffTime == null) {
                    this.cutOffTime = potentialCutoffTime.minus(delay).toEpochMilli();
                }
            }
        }

        if (iteratorWhoseHeadIsNext != null) {
            long cutoff = this.cutOffTime == null ? clock.instant().minus(delay).toEpochMilli() : this.cutOffTime;
            if (iteratorWhoseHeadIsNext.peek().event.effectiveTimestamp().isAfter(cutoff)) {
                underlying.clear();
                return null;
            }
        }
        return iteratorWhoseHeadIsNext;
    }
}

@@ -134,39 +134,8 @@
 tor;
-%0Aimport org.joda.time.DateTime;
 %0A%0Aim
@@ -209,16 +209,45 @@
 nstant;%0A
+import java.util.Comparator;%0A
 import j
@@ -616,16 +616,77 @@
 erlying;
+%0A    private final Comparator%3CEventInIdentifiedStream%3E order;
 %0A%0A    pr
@@ -904,24 +904,314 @@
 terators) %7B%0A
+        this(clock, delay, iterators, (a, b) -%3E a.event.effectiveTimestamp().compareTo(b.event.effectiveTimestamp()));%0A    %7D%0A%0A    public EventStoreStitchingIterator(Clock clock, Duration delay, List%3CIterator%3CEventInIdentifiedStream%3E%3E iterators, Comparator%3CEventInIdentifiedStream%3E order) %7B%0A
         this
@@ -1346,24 +1346,52 @@
 (toList());%0A
+        this.order = order;%0A
     %7D%0A%0A    @
@@ -2313,208 +2313,79 @@
     
-EventInIdentifiedStream candidate = eventStream.peek();%0A                DateTime candidateTimestamp = candidate.event.effectiveTimestamp();%0A                if (cutOffTime != null && candidateTimestamp
+if (cutOffTime != null && eventStream.peek().event.effectiveTimestamp()
 .isA
@@ -2508,36 +2508,42 @@
  %7C%7C 
-candidateTimestamp.isBefore(
+order.compare(eventStream.peek(), 
 iter
@@ -2572,36 +2572,13 @@
 ek()
-.event.effectiveTimestamp())
+) %3C 0
 ) %7B%0A

allow ordering strategy to be defined.