package com.atlassian.bitbucket.mesh.repository;

import org.apache.commons.lang3.mutable.MutableBoolean;
import org.apache.commons.lang3.mutable.MutableObject;
import org.junit.Before;
import org.junit.Test;

import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletionStage;

import static org.hamcrest.Matchers.instanceOf;
import static org.junit.Assert.*;

public class DrainableRequestTrackerTest {

    private int drainedCount;
    private DrainableRequestTracker tracker;

    @Before
    public void setUp() {
        drainedCount = 0;
        tracker = new DrainableRequestTracker();
    }

    @Test
    public void testCancelDrainingWhenDrained() {
        CompletionStage<Void> stage = tracker.asCompletionStage();
        MutableBoolean triggered = new MutableBoolean();
        MutableObject<Throwable> exception = new MutableObject<>();
        stage.handle((ignored, ex) -> {
            triggered.setTrue();
            exception.setValue(ex);
            return null;
        });

        //Start draining and let the tracker reach the drained state
        assertEquals(1, tracker.onRequestStarted());
        assertSame(stage, tracker.startDraining());
        assertEquals(0, tracker.onRequestEnded());
        assertTrue(triggered.booleanValue());
        assertNull(exception.getValue());

        //cancelDraining should replace the future, but not trigger further callbacks on the previous future
        assertEquals(0, tracker.cancelDraining());
        assertNull(exception.getValue());
        //The drainFuture should be replaced
        assertNotSame(stage, tracker.asCompletionStage());

        //For completeness, check that we can still start requests and subsequent draining works as expected
        stage = tracker.asCompletionStage();
        triggered.setFalse();
        exception.setValue(null);
        stage.handle((ignored, ex) -> {
            triggered.setTrue();
            exception.setValue(ex);
            return null;
        });

        assertEquals(1, tracker.onRequestStarted());
        assertEquals(2, tracker.onRequestStarted());
        assertSame(stage, tracker.startDraining()); // should still be the same (new) stage
        assertFalse(triggered.booleanValue());
        assertEquals(1, tracker.onRequestEnded());
        assertFalse(triggered.booleanValue());
        assertEquals(0, tracker.onRequestEnded());
        assertTrue(triggered.booleanValue());
    }

    @Test
    public void testCancelDrainingWhenDraining() {
        CompletionStage<Void> stage = tracker.asCompletionStage();
        MutableBoolean triggered = new MutableBoolean();
        MutableObject<Throwable> exception = new MutableObject<>();
        stage.handle((ignored, ex) -> {
            triggered.setTrue();
            exception.setValue(ex);
            return null;
        });

        assertEquals(1, tracker.onRequestStarted());
        assertSame(stage, tracker.startDraining());

        //cancelDraining should cancel the future
        assertEquals(1, tracker.cancelDraining());
        assertTrue(triggered.booleanValue());
        assertThat(exception.getValue(), instanceOf(CancellationException.class));
        //The drainFuture should be replaced
        assertNotSame(stage, tracker.asCompletionStage());

        //For completeness, check that we can still start requests and subsequent draining works as expected
        stage = tracker.asCompletionStage();
        triggered.setFalse();
        exception.setValue(null);
        stage.handle((ignored, ex) -> {
            triggered.setTrue();
            exception.setValue(ex);
            return null;
        });

        assertEquals(2, tracker.onRequestStarted());
        assertSame(stage, tracker.startDraining()); // should still be the same (new) stage
        assertFalse(triggered.booleanValue());
        assertEquals(1, tracker.onRequestEnded());
        assertFalse(triggered.booleanValue());
        assertEquals(0, tracker.onRequestEnded());
        assertTrue(triggered.booleanValue());
    }

    @Test
    public void testCancelDrainingWhenNotDraining() {
        CompletionStage<Void> stage = tracker.asCompletionStage();
        MutableBoolean triggered = new MutableBoolean();
        stage.handle((ignored, ex) -> {
            triggered.setTrue();
            return null;
        });

        assertEquals(1, tracker.onRequestStarted());

        //cancelDraining should no-op when we're not actually draining
        assertEquals(1, tracker.cancelDraining());
        assertFalse(triggered.booleanValue());

        //For completeness, check that we can still start requests and subsequent draining works as expected
        assertEquals(2, tracker.onRequestStarted());
        assertSame(stage, tracker.startDraining()); // should still be the same stage
        assertFalse(triggered.booleanValue());
        assertEquals(1, tracker.onRequestEnded());
        assertFalse(triggered.booleanValue());
        assertEquals(0, tracker.onRequestEnded());
        assertTrue(triggered.booleanValue());
    }

    @Test
    public void testInitialState() {
        //Verify that the tracker reports the correct values after creation
        assertEquals(0, tracker.getActiveRequestCount());
        assertEquals(0, drainedCount);
    }

    @Test
    public void testStartDrainingWhenNoRequestsAreActive() {
        tracker.startDraining().thenApply(ignored -> drainedCount++);

        //If the completion stage was completed, drainedCount would be 1
        assertEquals(1, drainedCount);

        //Verify that draining does not block new requests from being started
        assertEquals(1, tracker.onRequestStarted());
        //When the active request count reaches 0 again, the original CompletionStage should _not_ be triggered again
        assertEquals(0, tracker.onRequestEnded());
    }

    @Test
    public void testStartDrainingWhenRequestsActive() {
        //Mock 2 active requests
        assertEquals(1, tracker.onRequestStarted());
        assertEquals(2, tracker.onRequestStarted());
        assertEquals(2, tracker.getActiveRequestCount());

        //Start draining
        tracker.startDraining().thenApply(ignored -> drainedCount++);
        assertEquals(0, drainedCount);

        //Mock 1 active request ending
        assertEquals(1, tracker.onRequestEnded());
        assertEquals(1, tracker.getActiveRequestCount());
        assertEquals(0, drainedCount);

        //While we're draining, it's still allowed to start a new request
        assertEquals(2, tracker.onRequestStarted());
        assertEquals(2, tracker.getActiveRequestCount());
        assertEquals(0, drainedCount);

        //Mock all active requests ending
        assertEquals(1, tracker.onRequestEnded());
        assertEquals(0, tracker.onRequestEnded());
        assertEquals(1, drainedCount);

        //Verify that completed draining does not block new requests from being started
        assertEquals(1, tracker.onRequestStarted());
        //When the active request count reaches 0 again, the original CompletionStage should _not_ be triggered again
        assertEquals(0, tracker.onRequestEnded());

        //Verify that a new drain request works as expected
        assertEquals(1, tracker.onRequestStarted());
        tracker.startDraining().thenApply(ignored -> drainedCount++);
        assertEquals(1, drainedCount);
        assertEquals(0, tracker.onRequestEnded());
        assertEquals(2, drainedCount);
    }
}