package com.atlassian.stash.internal.maintenance.latch;

import com.atlassian.hazelcast.test.SpringManagedCluster;
import com.atlassian.johnson.Johnson;
import com.atlassian.johnson.JohnsonEventContainer;
import com.hazelcast.core.Cluster;
import com.hazelcast.core.HazelcastInstance;
import com.hazelcast.core.IExecutorService;
import org.junit.*;
import org.springframework.core.io.ClassPathResource;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.io.IOException;
import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkState;
import static org.junit.Assert.*;

public class ClusterableLatchTest {

    @ClassRule
    public static SpringManagedCluster cluster = new SpringManagedCluster.Builder().size(3).build();

    private SomeLatchableService node1Service;
    private SomeLatchableService node2Service;
    private SomeLatchableService node3Service;

    @BeforeClass
    public static void setupJohnson() throws IOException {
        Johnson.initialize(new ClassPathResource("minimal-johnson-config.xml").getURL().getPath());
    }

    @AfterClass
    public static void tearDownJohnson() {
        Johnson.terminate();
    }

    @Before
    public void setup() {
        cluster.reset();

        node1Service = createLatchableService(0);
        node2Service = createLatchableService(1);
        node3Service = createLatchableService(2);

        cluster.registerBeansForNode(0, node1Service);
        cluster.registerBeansForNode(1, node2Service);
        cluster.registerBeansForNode(2, node3Service);
    }

    @After
    public void clearAnyJohnsonEvents() {
        JohnsonEventContainer container = Johnson.getEventContainer();
        while (container.hasEvents()) {
            container.removeEvent(container.getEvents().iterator().next());
        }
    }

    @Test
    public void testAcquireCluster() {
        SomeLatchableService.SomeLatch latch1 = node1Service.acquireLatch(LatchMode.CLUSTER);

        // verify node1
        assertNotNull(latch1);
        assertEquals(LatchMode.CLUSTER, latch1.getMode());
        assertLatch(latch1, 1, 0, 0);

        // verify node2 and node 3
        assertLatch(node2Service.getCurrentLatch(), 1, 0, 0);
        assertLatch(node3Service.getCurrentLatch(), 1, 0, 0);
    }

    @Test(expected = IllegalStateException.class, timeout = 5000L)
    public void testAcquireClusterFailureUnlocks() throws InterruptedException {
        node2Service.failOnAcquire = true;
        node3Service.failOnUnlatch = true;

        try {
            node1Service.acquireLatch(LatchMode.CLUSTER);
            fail("node1Service.acquireLatch should have thrown an IllegalStateException");
        } finally {
            // node 1 and node 2 should be unlatched
            assertFalse(node1Service.isLatched());
            assertFalse(node2Service.isLatched());

            // node 3 should be Johnsoned, but it happens async so we may have to wait a bit for the Johnson task
            // to complete
            assertTrue(node3Service.isLatched());
            while (!Johnson.getEventContainer().hasEvents()) {
                Thread.sleep(25L);
            }
            assertTrue(Johnson.getEventContainer().hasEvents());
        }
    }

    @Test
    public void testDrainCluster() {
        SomeLatchableService.SomeLatch latch1 = node1Service.acquireLatch(LatchMode.CLUSTER);
        SomeLatchableService.SomeLatch latch2 = node2Service.getCurrentLatch();
        SomeLatchableService.SomeLatch latch3 = node3Service.getCurrentLatch();

        assertNotNull(latch2);
        assertNotNull(latch3);

        // both latches will fail to drain until SomeLatch.drained is set manually
        assertFalse(latch1.drain(25, TimeUnit.MILLISECONDS));
        assertFalse(latch2.drain(25, TimeUnit.MILLISECONDS));
        assertFalse(latch3.drain(25, TimeUnit.MILLISECONDS));

        // simulate one node having drained, both latches should still fail to drain
        latch1.drained = true;
        assertFalse(latch1.drain(25, TimeUnit.MILLISECONDS));
        assertFalse(latch2.drain(25, TimeUnit.MILLISECONDS));
        assertFalse(latch3.drain(25, TimeUnit.MILLISECONDS));

        // now simulate all nodes having drained
        latch2.drained = true;
        latch3.drained = true;
        assertTrue(latch1.drain(25, TimeUnit.MILLISECONDS));
        assertTrue(latch2.drain(25, TimeUnit.MILLISECONDS));
        assertTrue(latch3.drain(25, TimeUnit.MILLISECONDS));

        // latch1 had drained earlier. As an optimisation, it should not have been asked whether it had drained after
        // it had been discovered that it had already drained
        assertLatch(latch1, 1, 4, 0);
        assertLatch(latch2, 1, 7, 0);
        assertLatch(latch3, 1, 7, 0);
    }

    @Test
    public void testUnlatch() {
        node1Service.acquireLatch(LatchMode.CLUSTER);

        SomeLatchableService.SomeLatch latch3 = node3Service.getCurrentLatch();

        assertNotNull(latch3);

        // unlatch latch3 and verify that the other latches have unlatched as well
        latch3.unlatch();

        assertFalse(node1Service.isLatched());
        assertFalse(node2Service.isLatched());
        assertFalse(node3Service.isLatched());
    }

    private void assertLatch(SomeLatchableService.SomeLatch latch, int expectedAcquireCount, int expectedDrainCount,
                             int expectedUnlatchCount) {
        assertNotNull(latch);
        assertEquals(expectedAcquireCount, latch.acquiredCount);
        assertEquals(expectedDrainCount, latch.drainedCount);
        assertEquals(expectedUnlatchCount, latch.unlatchedCount);
    }

    private SomeLatchableService createLatchableService(int nodeIndex) {
        HazelcastInstance hazelcast = cluster.getNode(nodeIndex);
        return new SomeLatchableService(hazelcast.getCluster(), hazelcast.getExecutorService("test.executor"));
    }

    private static class SomeLatchableService implements LatchableService<Latch> {

        private final Cluster cluster;
        private final IExecutorService executorService;
        private volatile SomeLatch latch;

        boolean failOnAcquire = false;
        boolean failOnUnlatch = false;

        private SomeLatchableService(Cluster cluster, IExecutorService executorService) {
            this.cluster = cluster;
            this.executorService = executorService;
        }

        @Nonnull
        @Override
        public SomeLatch acquireLatch(@Nonnull LatchMode latchMode) {
            return acquireLatch(latchMode, null);
        }

        @Nonnull
        @Override
        public SomeLatch acquireLatch(@Nonnull LatchMode latchMode, String latchId) {
            checkState(latch == null, "already latched!");

            if (failOnAcquire) {
                throw new IllegalStateException("Simulating a failure in acquiring the latch");
            }

            SomeLatch someLatch = new SomeLatch(latchMode, cluster, executorService, getClass().getName());
            someLatch.acquire(latchId);
            this.latch = someLatch;
            return someLatch;
        }

        @Nullable
        @Override
        public SomeLatch getCurrentLatch() {
            return latch;
        }

        @Override
        public boolean isLatched() {
            return latch != null;
        }

        @Override
        public LatchState getState() {
            SomeLatch current = latch;

            return current == null ? LatchState.AVAILABLE :
                    current.drain(0, TimeUnit.MILLISECONDS) ? LatchState.DRAINED : LatchState.LATCHED;
        }

        class SomeLatch extends ClusterableLatch {

            volatile int acquiredCount;
            volatile int drainedCount;
            volatile int unlatchedCount;

            volatile boolean drained;

            protected SomeLatch(LatchMode mode, Cluster cluster, IExecutorService executor, String latchServiceBeanName) {
                super(mode, cluster, executor, latchServiceBeanName);
            }

            @Override
            protected void acquireLocally() {
                acquiredCount++;
            }

            @Override
            protected boolean drainLocally(long timeout, @Nonnull TimeUnit timeUnit, boolean force) {
                drainedCount++;
                return drained;
            }

            @Override
            protected void unlatchLocally() {
                unlatchedCount++;
                if (failOnUnlatch) {
                    throw new IllegalStateException("Simulating failure to unlatch");
                }
                latch = null;
            }
        }
    }
}