/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.streaming.runtime.tasks;

import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.concurrent.Delayed;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.util.Preconditions;

/**
 * This is a {@link ProcessingTimeService} used <b>strictly for testing</b> the
 * processing time functionality.
 * */
public class TestProcessingTimeService extends ProcessingTimeService {

	private volatile long currentTime = Long.MIN_VALUE;

	private volatile boolean isTerminated;
	private volatile boolean isQuiesced;

	// sorts the timers by timestamp so that they are processed in the correct order.
	private final PriorityQueue<Tuple2<Long, CallbackTask>> priorityQueue;

	public TestProcessingTimeService() {
		this.priorityQueue = new PriorityQueue<>(16, new Comparator<Tuple2<Long, CallbackTask>>() {
			@Override
			public int compare(Tuple2<Long, CallbackTask> o1, Tuple2<Long, CallbackTask> o2) {
				return Long.compare(o1.f0, o2.f0);
			}
		});
	}

	public void setCurrentTime(long timestamp) throws Exception {
		this.currentTime = timestamp;

		if (!isQuiesced) {
			while (!priorityQueue.isEmpty() && currentTime >= priorityQueue.peek().f0) {
				Tuple2<Long, CallbackTask> entry = priorityQueue.poll();

				CallbackTask callbackTask = entry.f1;

				if (!callbackTask.isDone()) {
					callbackTask.onProcessingTime(entry.f0);

					if (callbackTask instanceof PeriodicCallbackTask) {
						priorityQueue.offer(Tuple2.of(((PeriodicCallbackTask) callbackTask).nextTimestamp(entry.f0), callbackTask));
					}
				}
			}
		}
	}

	@Override
	public long getCurrentProcessingTime() {
		return currentTime;
	}

	@Override
	public ScheduledFuture<?> registerTimer(long timestamp, ProcessingTimeCallback target) {
		if (isTerminated) {
			throw new IllegalStateException("terminated");
		}
		if (isQuiesced) {
			return new CallbackTask(null);
		}

		CallbackTask callbackTask = new CallbackTask(target);

		priorityQueue.offer(Tuple2.of(timestamp, callbackTask));

		return callbackTask;
	}

	@Override
	public ScheduledFuture<?> scheduleAtFixedRate(ProcessingTimeCallback callback, long initialDelay, long period) {
		if (isTerminated) {
			throw new IllegalStateException("terminated");
		}
		if (isQuiesced) {
			return new CallbackTask(null);
		}

		PeriodicCallbackTask periodicCallbackTask = new PeriodicCallbackTask(callback, period);

		priorityQueue.offer(Tuple2.<Long, CallbackTask>of(currentTime + initialDelay, periodicCallbackTask));

		return periodicCallbackTask;
	}

	@Override
	public boolean isTerminated() {
		return isTerminated;
	}

	@Override
	public void quiesceAndAwaitPending() {
		if (!isTerminated) {
			isQuiesced = true;
			priorityQueue.clear();
		}
	}

	@Override
	public void shutdownService() {
		this.isTerminated = true;
	}

	public int getNumActiveTimers() {
		int count = 0;

		for (Tuple2<Long, CallbackTask> entry : priorityQueue) {
			if (!entry.f1.isDone()) {
				count++;
			}
		}

		return count;
	}

	public Set<Long> getActiveTimerTimestamps() {
		Set<Long> actualTimestamps = new HashSet<>();

		for (Tuple2<Long, CallbackTask> entry : priorityQueue) {
			if (!entry.f1.isDone()) {
				actualTimestamps.add(entry.f0);
			}
		}

		return actualTimestamps;
	}

	// ------------------------------------------------------------------------

	private static class CallbackTask implements ScheduledFuture<Object> {

		protected final ProcessingTimeCallback processingTimeCallback;

		private AtomicReference<CallbackTaskState> state = new AtomicReference<>(CallbackTaskState.CREATED);

		private CallbackTask(ProcessingTimeCallback processingTimeCallback) {
			this.processingTimeCallback = processingTimeCallback;
		}

		public void onProcessingTime(long timestamp) throws Exception {
			processingTimeCallback.onProcessingTime(timestamp);

			state.compareAndSet(CallbackTaskState.CREATED, CallbackTaskState.DONE);
		}

		@Override
		public long getDelay(TimeUnit unit) {
			throw new UnsupportedOperationException();
		}

		@Override
		public int compareTo(Delayed o) {
			throw new UnsupportedOperationException();
		}

		@Override
		public boolean cancel(boolean mayInterruptIfRunning) {
			return state.compareAndSet(CallbackTaskState.CREATED, CallbackTaskState.CANCELLED);
		}

		@Override
		public boolean isCancelled() {
			return state.get() == CallbackTaskState.CANCELLED;
		}

		@Override
		public boolean isDone() {
			return state.get() != CallbackTaskState.CREATED;
		}

		@Override
		public Object get() throws InterruptedException, ExecutionException {
			throw new UnsupportedOperationException();
		}

		@Override
		public Object get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
			throw new UnsupportedOperationException();
		}

		enum CallbackTaskState {
			CREATED,
			CANCELLED,
			DONE
		}
	}

	private static class PeriodicCallbackTask extends CallbackTask {

		private final long period;

		private PeriodicCallbackTask(ProcessingTimeCallback processingTimeCallback, long period) {
			super(processingTimeCallback);
			Preconditions.checkArgument(period > 0L, "The period must be greater than 0.");

			this.period = period;
		}

		@Override
		public void onProcessingTime(long timestamp) throws Exception {
			processingTimeCallback.onProcessingTime(timestamp);
		}

		public long nextTimestamp(long currentTimestamp) {
			return currentTimestamp + period;
		}
	}
}