package com.atlassian.bitbucket.internal.mirroring.mirror.sync;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.ListenableFutureTask;
import com.google.common.util.concurrent.MoreExecutors;

import javax.annotation.Nonnull;
import java.util.Optional;
import java.util.concurrent.*;

import static java.util.Optional.empty;
import static java.util.Optional.of;

/**
 * An executor that does not allow multiple instances of the same task (identified by taskId) to run concurrently.
 * If a task is submitted while it is already running, it is prevented from executing until the running task completes.
 * If a task is submitted n times while the task is already running, it will run only once after the task completes.
 */
public class DedupingExecutor {

    private final ExecutorService delegate;
    private final ConcurrentMap<Object, Optional<FutureTask<?>>> submittedWhileExecuting;

    public DedupingExecutor(ExecutorService delegate) {
        this.delegate = delegate;

        submittedWhileExecuting = Maps.newConcurrentMap();
    }

    public Future<?> submit(@Nonnull Object taskId, Runnable task) {
        // Task is being submitted. We need to ensure that the same task does not execute more than once concurrently
        //
        // If the task is not already running, add (taskId, empty()) to the map. This marks the task as currently running
        // and states that there is no subsequent task that needs to run (yet).
        //
        // If the task is already running, update (taskId, empty()) --> (taskId, of(FutureTask(task))) to mark it as
        // needing to run the provided task next
        //
        // If the task is already running and a task has already been queued up, ignore the provided task and return
        // the Future to the queued up task
        Optional<FutureTask<?>> next = submittedWhileExecuting.merge(taskId, Optional.<FutureTask<?>>empty(),
                (oldValue, newValue) -> {
                    if (oldValue.isPresent()) {
                        // another task has already been queued up, keep that
                        return oldValue;
                    }
                    // no task has been queued up yet, set the submitted task as the next task to run
                    return of(resubmittingTask(taskId, task));
        });

        if (next.isPresent()) {
            // a task has already been queued up. Return that
            return next.get();
        }

        // no task was running, execute the task
        try {
            return delegate.submit(resubmittingTask(taskId, task));
        } catch (RejectedExecutionException e) {
            // unmark the task as running
            submittedWhileExecuting.remove(taskId);
            throw e;
        }
    }

    private FutureTask<?> resubmittingTask(Object taskId, Runnable task) {
        ListenableFutureTask<?> futureTask = ListenableFutureTask.create(task, null);
        futureTask.addListener(new MaybeSubmitNextTask(taskId), MoreExecutors.directExecutor());
        return futureTask;
    }

    private class MaybeSubmitNextTask implements Runnable {
        private final Object taskId;

        private MaybeSubmitNextTask(Object taskId) {
            this.taskId = taskId;
        }

        @Override
        public void run() {
            // loop to handle concurrent modifications of the submittedWhileExecuting map
            for (;;) {
                Optional<FutureTask<?>> nextTask = submittedWhileExecuting.get(taskId);
                if (nextTask == null) {
                    // nothing to do
                    return;
                }
                if (!nextTask.isPresent()) {
                    // no nextTask was scheduled, clean up and return
                    if (submittedWhileExecuting.remove(taskId, nextTask)) {
                        return;
                    }
                }

                if (submittedWhileExecuting.replace(taskId, nextTask, empty())) {
                    try {
                        delegate.submit(nextTask.get());
                        return;
                    } catch (RejectedExecutionException e) {
                        // unmark the task as running
                        submittedWhileExecuting.remove(taskId);
                        throw e;
                    }
                }
            }
        }
    }
}
