package com.atlassian.bitbucket.mesh.util;

import javax.annotation.Nonnull;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Serializes operations on a per key basis.
 *
 * @param <K> the type of the key
 */
public class SerializedOperation<K> {

    private final ConcurrentMap<K, LockEntry> perKeyLocks;

    public SerializedOperation() {
        perKeyLocks = new ConcurrentHashMap<>();
    }

    /**
     * Submit an operation, encapsulated by a {@link java.util.concurrent.Callable}, and keyed by an
     * Object {@literal <K>}, such that operations with the same key will not execute concurrently, but rather
     * execution will be serialised.
     * <p>
     * No guarantees are made as to the order of execution of submitted operations, however the implementation
     * does favor granting access to the longest-waiting thread to avoid starvation.
     *
     * @param key the key, this should be an immutable object that correctly implements
     *            {@link java.lang.Object#hashCode()} and {@link java.lang.Object#equals(Object)}
     * @param operation is the operation to execute
     * @param <R> the type the operation returns
     * @return result of the operation
     * @throws Exception if the callable generated an exception
     */
    @Nonnull
    public <R> R runOperation(@Nonnull K key, @Nonnull Callable<R> operation) throws Exception {
        LockEntry perKeyLock = perKeyLocks.compute(key, (ignored, value) ->
                value == null ? new LockEntry() : value.incrementReferenceCount());

        try {
            perKeyLock.lock();
            return operation.call();
        } finally {
            perKeyLock.unlock();
            perKeyLocks.compute(key, (ignored, value) -> {
                if (value == null) {
                    // If this happens it is an implementation error, not a caller error
                    throw new IllegalStateException("Referenced lock missing from map");
                }
                return value.decrementAndGetReferenceCount() == 0 ? null : value;
            });
        }
    }

    private static class LockEntry {

        private final AtomicInteger referenceCount;
        private final ReentrantLock lock;

        LockEntry() {
            referenceCount = new AtomicInteger(1);
            lock = new ReentrantLock(true);
        }

        int decrementAndGetReferenceCount() {
            return referenceCount.decrementAndGet();
        }

        LockEntry incrementReferenceCount() {
            referenceCount.incrementAndGet();
            return this;
        }

        void lock() {
            lock.lock();
        }

        void unlock() {
            lock.unlock();
        }
    }
}
