package com.atlassian.bitbucket.mesh.repository;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;

class DrainableRequestTracker {

    private final Object lock;

    private int activeRequests;
    private CompletableFuture<Void> drainFuture;
    private boolean draining;

    DrainableRequestTracker() {
        this(0);
    }

    DrainableRequestTracker(int count) {
        activeRequests = count;
        lock = new Object();
        drainFuture = new CompletableFuture<>();
    }

    CompletionStage<Void> asCompletionStage() {
        synchronized (lock) {
            return drainFuture;
        }
    }

    int cancelDraining() {
        synchronized (lock) {
            if (draining) {
                draining = false;
                //Cancel the current future to unlock any threads currently waiting for this tracker to drain
                drainFuture.cancel(false);
                //Replace the future with a fresh future
                drainFuture = new CompletableFuture<>();
            }
            return activeRequests;
        }
    }

    int getActiveRequestCount() {
        synchronized (lock) {
            return activeRequests;
        }
    }

    CompletionStage<Void> startDraining() {
        synchronized (lock) {
            draining = true;
            if (activeRequests == 0) {
                //We're drained!
                drainFuture.complete(null);
            }
            return drainFuture;
        }
    }

    int onRequestEnded() {
        synchronized (lock) {
            if (--activeRequests == 0 && draining) {
                //We've completed draining!
                drainFuture.complete(null);
                //Replace the drainFuture to prepare for the next drain request
                draining = false;
                drainFuture = new CompletableFuture<>();
            }
            return activeRequests;
        }
    }

    int onRequestStarted() {
        synchronized (lock) {
            return ++activeRequests;
        }
    }
}
