package org.codefx.demo.java9.api.reactive_streams;

import java.util.Collections;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Flow.Publisher;
import java.util.concurrent.Flow.Subscriber;
import java.util.concurrent.Flow.Subscription;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import static java.util.concurrent.Executors.newSingleThreadExecutor;

/**
 * A publisher that produces a potentially infinite stream of consecutive positive integers.
 * <br>
 * The publisher makes a best effort to start each subscription's (partial) series
 * with the smallest value that has not yet been requested by all active subscribers.
 * (This item could be seen to be the oldest that was not yet processed by all existing
 * subscriptions.)
 * <br>
 * Once the last subscription is cancelled, the publisher will cease to accept new
 * subscriptions - {@link #waitUntilTerminated()} will terminate soon after, once
 * all remaining items were sent out.
 */
public class IncrementingPublisher implements Publisher<Integer> {

	private final ExecutorService executor = Executors.newFixedThreadPool(4);
	private final Set<Sub> subscriptions = Collections.newSetFromMap(new ConcurrentHashMap<>());
	private final AtomicInteger subscriptionCount = new AtomicInteger();
	private final CompletableFuture<Void> terminated = new CompletableFuture<>();

	@Override
	public void subscribe(Subscriber<? super Integer> subscriber) {
		Sub subscription = createNewSubscriptionFor(subscriber);
		registerSubscription(subscription);
		subscriber.onSubscribe(subscription);
	}

	private Sub createNewSubscriptionFor(Subscriber<? super Integer> subscriber) {
		int startValue = subscriptions.stream()
				.mapToInt(sub -> sub.nextValue.get())
				.min()
				.orElse(0);
		return new Sub(subscriber, startValue);
	}

	private void registerSubscription(Sub subscription) {
		subscriptions.add(subscription);
		subscriptionCount.incrementAndGet();
	}

	private boolean unregisterSubscriptionAndCheckIfLast(Sub subscription) {
		subscriptions.remove(subscription);
		return subscriptionCount.decrementAndGet() == 0;
	}

	private void shutdown() {
		System.out.println("Shutting down executor service...");
		executor.shutdown();
		newSingleThreadExecutor().submit(() -> {
			try {
				executor.awaitTermination(0, TimeUnit.SECONDS);
			} catch (InterruptedException ex) {
				// if waiting gets interrupted, we simply declare the publisher
				// to be terminated
			}
			System.out.println("Shutdown complete.");
			terminated.complete(null);
		});
	}

	public void waitUntilTerminated() throws InterruptedException {
		try {
			terminated.get();
		} catch (ExecutionException ex) {
			// even if something went wrong - the computation is terminated all the same
			System.out.println(ex);
		}
	}

	private class Sub implements Subscription {

		private final Subscriber<? super Integer> subscriber;
		private final AtomicInteger nextValue;
		private final AtomicBoolean canceled;

		public Sub(Subscriber<? super Integer> subscriber, int startValue) {
			this.subscriber = subscriber;
			this.nextValue = new AtomicInteger(startValue);
			this.canceled = new AtomicBoolean(false);
		}

		@Override
		public void request(long n) {
			if (canceled.get())
				return;

			if (n < 0)
				reportIllegalArgument();
			else
				publishItems(n);
		}

		private void reportIllegalArgument() {
			executor.execute(() -> subscriber.onError(new IllegalArgumentException()));
		}

		private void publishItems(long n) {
			for (long i = n; i > 0; i--)
				executor.execute(() -> subscriber.onNext(nextValue.getAndIncrement()));
		}

		@Override
		public void cancel() {
			canceled.set(true);
			// we do not cancel already requested items;
			// instead we unregister check whether we want to shut down
			boolean wasLast = unregisterSubscriptionAndCheckIfLast(this);
			if (wasLast)
				shutdown();
		}
	}
}