package com.atlassian.bitbucket.mesh.util;

import com.google.common.annotations.VisibleForTesting;

import javax.annotation.Nonnull;
import java.time.Duration;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.locks.ReentrantLock;

import static java.util.Objects.requireNonNull;

/**
 * A utility class for restricting how often a given operation is run. Each key is gated separately, allowing each
 * to run or not run based on when an operation was last run for that specific key.
 *
 * @param <T> the key type
 */
public class Gate<T> {

    @VisibleForTesting
    static final int CLEANUP_THRESHOLD = 1500;

    private final long intervalNanos;
    private final ReentrantLock lock;
    private final ConcurrentMap<T, Long> timestampsByKey;

    /**
     * Constructs a new {@code Gate} which will allow operations for each key to run at most once within the
     * specified interval.
     *
     * @param interval the gate interval
     */
    public Gate(@Nonnull Duration interval) {
        intervalNanos = requireNonNull(interval, "interval").toNanos();
        lock = new ReentrantLock();
        timestampsByKey = new ConcurrentHashMap<>();
    }

    /**
     * Runs the provided operation if no other operation has run for the same key within the interval configured
     * when the {@code Gate} was constructed.
     *
     * @param key       the key, used to differentiate intervals
     * @param operation the operation to perform
     * @return {@code true} if the provided operation ran; otherwise, {@code false} if another operation had run
     *         too recently
     */
    public boolean maybeRun(@Nonnull T key, @Nonnull Runnable operation) {
        requireNonNull(key, "key"); //ConcurrentMaps don't allow null keys
        requireNonNull(operation, "operation");

        long now = System.nanoTime();
        long timestamp = timestampsByKey.compute(key,
                (k, lastRun) -> (lastRun == null || lastRun + intervalNanos <= now) ? now : lastRun);

        boolean hasRun;
        if (timestamp == now) {
            operation.run();
            hasRun = true;
        } else {
            hasRun = false;
        }

        if (CLEANUP_THRESHOLD < timestampsByKey.size()) {
            cleanup(now);
        }

        return hasRun;
    }

    /**
     * @return the number of keys for which an interval is tracked, some of which may be expired
     */
    @VisibleForTesting
    protected int getMapSize() {
        return timestampsByKey.size();
    }

    /**
     * Removes expired entries (entries older than the configured interval) from the map.
     *
     * @param now the current timestamp in nanoseconds
     */
    private void cleanup(long now) {
        if (lock.tryLock()) {
            try {
                timestampsByKey.entrySet().removeIf(entry -> entry.getValue() + intervalNanos <= now);
            } finally {
                lock.unlock();
            }
        }
    }
}