package rxbroadcast;

import rxbroadcast.time.Clock;
import rxbroadcast.time.LamportClock;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.function.Consumer;

public final class SingleSourceFifoOrder<T> implements BroadcastOrder<Timestamped<T>, T> {
    @SuppressWarnings("WeakerAccess")
    public static final boolean DROP_LATE = true;

    private final Clock clock = new LamportClock();

    private final Map<Sender, SortedSet<Timestamped<T>>> pendingQueues;

    private final Map<Sender, Long> expectedTimestamps;

    private final boolean dropLateMessages;

    public SingleSourceFifoOrder() {
        this(!DROP_LATE);
    }

    @SuppressWarnings("WeakerAccess")
    public SingleSourceFifoOrder(final boolean dropLateMessages) {
        this.pendingQueues = new HashMap<>();
        this.expectedTimestamps = new HashMap<>();
        this.dropLateMessages = dropLateMessages;
    }

    @Override
    public Timestamped<T> prepare(final T value) {
        return clock.tick(time -> new Timestamped<>(time, value));
    }

    @Override
    public void receive(final Sender sender, final Consumer<T> consumer, final Timestamped<T> value) {
        expectedTimestamps.putIfAbsent(sender, 1L);

        if (dropLateMessages) {
            if (Long.compareUnsigned(value.timestamp, expectedTimestamps.get(sender)) >= 0) {
                consumer.accept(value.value);
                expectedTimestamps.compute(sender, (k, v) -> value.timestamp + 1);
            }

            return;
        }

        final SortedSet<Timestamped<T>> queue = pendingQueues.computeIfAbsent(sender, k -> new TreeSet<>());

        if (value.timestamp == expectedTimestamps.get(sender)) {
            consumer.accept(value.value);
            expectedTimestamps.compute(sender, (k, v) -> value.timestamp + 1);
        } else {
            queue.add(value);
        }

        final Iterator<Timestamped<T>> iterator = queue.iterator();
        while (iterator.hasNext()) {
            final Timestamped<T> tv = iterator.next();
            if (tv.timestamp < expectedTimestamps.get(sender)) {
                iterator.remove();
                continue;
            }
            if (tv.timestamp > expectedTimestamps.get(sender)) {
                break;
            }

            consumer.accept(tv.value);
            expectedTimestamps.compute(sender, (k, v) -> tv.timestamp + 1);
            iterator.remove();
        }
    }
}

@@ -1712,215 +1712,8 @@
 );%0A%0A
-        if (value.timestamp == expectedTimestamps.get(sender)) %7B%0A            consumer.accept(value.value);%0A            expectedTimestamps.compute(sender, (k, v) -%3E value.timestamp + 1);%0A        %7D else %7B%0A    
     
@@ -1733,26 +1733,16 @@
 (value);
-%0A        %7D
 %0A%0A      

Remove unneeded timestamp check in SSF order