package com.atlassian.bitbucket.mesh.io;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.LockSupport;

/**
 * Utility class that can be used to apply throttling for I/O operations.
 * <p>
 * Note: this class is not thread-safe in that only one thread is expected to call {@link #maybeThrottle} at a time.
 * Other threads are allowed to call {@link #cancel()} though to release a throttled thread.
 */
public class IoThrottle {

    private static final Logger log = LoggerFactory.getLogger(IoThrottle.class);

    private final long bytesPerSecond;
    private final long nanosPerByte;
    private final long startTimestamp;
    private final AtomicLong totalBytesProcessed;

    private volatile boolean canceled;
    private volatile Thread parkedThread;

    private IoThrottle(long bytesPerSecond) {
        this.bytesPerSecond = bytesPerSecond;

        nanosPerByte = bytesPerSecond > 0 ? TimeUnit.SECONDS.toNanos(1) / bytesPerSecond : -1L;
        startTimestamp = System.nanoTime();
        totalBytesProcessed = new AtomicLong();
    }

    @Nonnull
    public static IoThrottle start(long bytesPerSecond) {
        return new IoThrottle(bytesPerSecond);
    }

    /**
     * Releases the parked thread, if it is currently parked waiting (throttling)
     */
    public void cancel() {
        canceled = true;

        Thread thread = parkedThread;
        if (thread != null) {
            LockSupport.unpark(thread);
        }
    }

    public void maybeThrottle(long bytesProcessed) {
        if (canceled) {
            return;
        }
        if (nanosPerByte > 0) {
            long bytes = totalBytesProcessed.addAndGet(bytesProcessed);
            long nextEmitTimestamp = startTimestamp + (nanosPerByte * bytes);
            long pauseNanos;
            while (!canceled && (pauseNanos = nextEmitTimestamp - System.nanoTime()) > 0) {
                if (log.isTraceEnabled()) {
                    long sleepMs = TimeUnit.NANOSECONDS.toMillis(pauseNanos);
                    long timeSinceStart = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startTimestamp);
                    log.trace("Throttling file processing; parking for {}ms before sending next message " +
                            "(current rate = {} bytes/s, target rate = {} bytes/s)",
                            sleepMs, (bytes * 1000.0d / timeSinceStart), bytesPerSecond);
                }
                parkedThread = Thread.currentThread();
                LockSupport.parkNanos(pauseNanos);
            }
        }
    }
}
