package com.atlassian.stash.internal.maintenance;

import com.atlassian.event.api.EventPublisher;
import com.atlassian.hazelcast.test.SpringManagedCluster;
import com.atlassian.johnson.Johnson;
import com.atlassian.security.random.SecureTokenGenerator;
import com.atlassian.bitbucket.AbstractServiceTest;
import com.atlassian.bitbucket.cluster.ClusterInformation;
import com.atlassian.bitbucket.cluster.ClusterNode;
import com.atlassian.bitbucket.cluster.ClusterService;
import com.atlassian.bitbucket.event.cluster.ClusterNodeAddedEvent;
import com.atlassian.bitbucket.AuthorisationException;
import com.atlassian.bitbucket.i18n.I18nService;
import com.atlassian.bitbucket.i18n.KeyedMessage;
import com.atlassian.bitbucket.i18n.SimpleI18nService;
import com.atlassian.stash.internal.backup.CanceledBackupException;
import com.atlassian.stash.internal.db.DatabaseManager;
import com.atlassian.stash.internal.hazelcast.HazelcastConstants;
import com.atlassian.stash.internal.maintenance.latch.LatchState;
import com.atlassian.stash.internal.scm.InternalScmService;
import com.atlassian.bitbucket.request.RequestContext;
import com.atlassian.bitbucket.request.RequestManager;
import com.atlassian.bitbucket.user.ApplicationUserVisitor;
import com.atlassian.bitbucket.auth.AuthenticationContext;
import com.atlassian.bitbucket.user.ApplicationUser;
import com.atlassian.bitbucket.user.UserType;
import com.atlassian.bitbucket.util.Progress;
import com.atlassian.bitbucket.util.ProgressImpl;
import com.google.common.base.Throwables;
import com.hazelcast.core.HazelcastInstance;
import com.hazelcast.core.IAtomicLong;
import com.hazelcast.core.IAtomicReference;
import com.hazelcast.core.IExecutorService;
import org.junit.*;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.Spy;
import org.mockito.runners.MockitoJUnitRunner;
import org.springframework.core.io.ClassPathResource;

import javax.annotation.Nonnull;
import java.io.IOException;
import java.io.Serializable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import static org.hamcrest.Matchers.equalTo;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultMaintenanceServiceTest extends AbstractServiceTest {

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

    @Mock
    private AuthenticationContext authenticationContext;
    private ClusterService[] clusterServices;
    private int counter;
    @Mock
    private DatabaseManager databaseManager;
    @Mock
    private EventPublisher eventPublisher;
    private ScheduledExecutorService executorService;
    @Spy
    private I18nService i18nService = new SimpleI18nService();
    @Spy
    private MaintenanceTaskStatusSupplier latestTask = new LocalMaintenanceTaskStatusSupplier();
    @Mock
    private RequestManager requestManager;
    @Mock
    private InternalScmService scmService;
    private DefaultMaintenanceService service1;
    private DefaultMaintenanceService service2;
    @Mock
    private SecureTokenGenerator tokenGenerator;

    public DefaultMaintenanceServiceTest() {
        super(DefaultMaintenanceService.class, MaintenanceService.class);
    }

    @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 setupMocks() {
        cluster.reset();
        ++counter;

        Assume.assumeThat("nodes haven't formed a cluster!", cluster.getNode(0).getCluster().getMembers().size(), equalTo(2));
        Assume.assumeThat("nodes haven't formed a cluster!", cluster.getNode(1).getCluster().getMembers().size(), equalTo(2));

        executorService = Executors.newSingleThreadScheduledExecutor();
        clusterServices = new ClusterService[2];

        service1 = createMaintenanceService(0);
        service2 = createMaintenanceService(1);
    }

    @After
    public void teardown() {
        executorService.shutdownNow();
    }

    @Test
    public void testGetLockContextOutsideMaintenance() {
        assertNull(service1.getLock());
    }

    @Test
    public void testGetStatus() {
        when(databaseManager.getState()).thenReturn(LatchState.LATCHED);
        when(scmService.getState()).thenReturn(LatchState.DRAINED);

        MaintenanceStatus status = service1.getStatus();
        assertNotNull(status);
        assertEquals(LatchState.LATCHED, status.getDatabaseState());
        assertEquals(LatchState.DRAINED, status.getScmState());
    }

    @Test
    public void testGetTaskContextOutsideMaintenance() {
        assertNull(service1.getRunningTask());
    }

    @Test
    public void testJoinClusterWhileClusterIsLocking() throws InterruptedException {
        mockAuthenticatedAs("user-join-cluster-while-cluster-is-locking", 6);

        // there is a race condition where the cluster is being in maintenance mode just as a node is joining the
        // cluster. The newly joined/joining node may not receive the IExecutor command that locks the node, leaving
        // the node unlocked while the rest of the cluster is locked.

        // set a the clusterLock
        IAtomicReference<ClusterMaintenanceLock> clusterLock = cluster.getNode(0).getAtomicReference("cluster-lock-" + counter);
        clusterLock.set(new ClusterMaintenanceLock(authenticationContext.getCurrentUser(), "pre-defined-token", null, null, null));

        // verify that the node isn't locked yet
        assertNull(service1.getNodeLock());

        // notify that a node has been added to the cluster
        service1.onNodeAdded(mock(ClusterNodeAddedEvent.class));

        // sleep for a very short time to allow node 1 to lock itself on the executor thread
        Thread.sleep(50L);

        MaintenanceLock lock = service1.getNodeLock();
        assertNotNull(lock);
        assertEquals("pre-defined-token", lock.getUnlockToken());
    }

    @Test(expected = AuthorisationException.class)
    public void testLockAnonymously() {
        service1.lock();
    }

    @Test
    public void testLockAndUnlockDifferentNodes() {
        final String token = "token-lock-and-unlock-different-nodes";

        mockAuthenticatedAs("user-lock-and-unlock-different-nodes", 7);
        when(tokenGenerator.generateToken()).thenReturn(token, "other");

        MaintenanceLock lock = service1.lock();
        assertNotNull(lock);

        // verify that the system is locked for maintenance on the other node as well
        MaintenanceLock retrieved = service2.getLock();
        assertNotNull(retrieved);
        assertEquals(retrieved, service2.getLock());
        assertEquals(token, retrieved.getUnlockToken());
        assertSame(7, retrieved.getOwner().getId());

        retrieved.unlock(token);

        // verify that the system is no longer locked
        assertNull(service1.getLock());
        assertNull(service2.getLock());
    }

    @Test
    public void testLockAndUnlockSameNode() {
        final String token = "token-lock-and-unlock-same-node";

        mockAuthenticatedAs("user-lock-and-unlock-same-node", 3);
        when(tokenGenerator.generateToken()).thenReturn(token);

        MaintenanceLock lock = service1.lock();

        // verify that the system is locked for maintenance
        assertNotNull(lock);
        assertNotNull(service1.getLock());
        assertNotNull(service2.getLock());
        assertEquals(token, lock.getUnlockToken());
        assertSame(3, lock.getOwner().getId());

        lock.unlock(token);

        // verify that the system is no longer locked
        assertNull(service1.getLock());
        assertNull(service2.getLock());
    }

    @Test(expected = LockedMaintenanceException.class)
    public void testLockWhileLocked() {
        mockAuthenticatedAs("user-lock-while-locked", 1);
        when(tokenGenerator.generateToken()).thenReturn("token-lock-while-locked");

        service1.lock();
        service1.lock();
    }

    // test what happens when system is locked for maintenance while the one node is already locked. This is an
    // inconsistent state, possibly due to a split-brain. It shouldn't be possible because of the join checks, but
    // still worth testing
    @Test
    public void testLockWhileSingleNodeLocked() {
        mockAuthenticatedAs("user-lock-while-single-node-locked", 1);
        when(tokenGenerator.generateToken()).thenReturn("token-lock-while-single-node-locked");

        // sanity check that the clusterLock is not set
        IAtomicReference<ClusterMaintenanceLock> clusterLock = cluster.getNode(0).getAtomicReference("cluster-lock-" + counter);
        Assume.assumeTrue("clusterlock should not be present!", clusterLock.get() == null);

        // lock node 1 locally - no cluster lock
        service1.lockNode(new DefaultMaintenanceLock(eventPublisher, i18nService, authenticationContext.getCurrentUser(), "node1-local-token"));
        assertNotNull("service 1 was not locked", service1.getLock());
        assertNotNull("service 1 was not locked locally", service1.getNodeLock());

        try {
            service2.lock();
            fail(LockFailedMaintenanceException.class.getName() + " expected");
        } catch (LockFailedMaintenanceException e) {
            // verify that node 2 is not locked
            assertNull(service2.getLock());
            assertNull(service2.getNodeLock());
        }
    }

    @Test(expected = IllegalStateException.class, timeout = 5000L)
    public void testStartFromBackgroundThread() {
        ClusterMaintenanceTask task = mock(ClusterMaintenanceTask.class);
        service1.start(task, MaintenanceType.MIGRATION);
    }

    @Test(timeout = 5000L)
    public void testStartClusterableTaskInCluster() throws InterruptedException {
        doTestStart(new ClusterMaintenanceTask("some progress", 2));
    }

    @Test(timeout = 5000L)
    public void testStartClusterableTaskInStandalone() throws InterruptedException {
        mockStandaloneSetup();

        // starting a ClusterableTask in a standalone install should be possible
        doTestStart(new ClusterMaintenanceTask("some progress", 2));
    }

    @Test(expected = UnsupportedMaintenanceException.class, timeout = 5000L)
    public void testStartStandaloneTaskInCluster() throws InterruptedException {
        doTestStart(new StandaloneMaintenanceTask("some progress", 2));
    }

    @Test(timeout = 5000L)
    public void testStartStandaloneTaskInStandalone() throws InterruptedException {
        mockStandaloneSetup();

        doTestStart(new StandaloneMaintenanceTask("some progress", 2));
    }

    @Test(expected = IllegalStateException.class)
    public void testStartWhileTaskIsRunning() {
        final String sessionId = "ssssssessssionId";
        final String token = "cancel-me";

        RequestContext requestContext = mock(RequestContext.class);
        when(requestContext.getSessionId()).thenReturn(sessionId);
        when(requestManager.getRequestContext()).thenReturn(requestContext);
        when(tokenGenerator.generateToken()).thenReturn(token);

        service1.start(new ClusterMaintenanceTask("progress", 1), MaintenanceType.MIGRATION);
        service1.start(new ClusterMaintenanceTask("progress task2", 1), MaintenanceType.MIGRATION);
    }

    @Test
    public void testRemotelyRunningTaskCanBeCanceled() throws Exception {
        mockAuthenticatedAs("user-remotely-running-task-can-be-cancelled", 12);
        ClusterMaintenanceTask task = new ClusterMaintenanceTask("progress", 5);
        when(tokenGenerator.generateToken()).thenReturn("cancel-token-remotely-running-task-can-be-cancelled");

        service1.start(task, MaintenanceType.BACKUP);

        // retrieve the current monitor
        MaintenanceTaskMonitor monitor1 = service1.getRunningTask();
        MaintenanceTaskMonitor monitor2 = service2.getRunningTask();

        // because it's running remotely, monitor should be non-null
        assertNotNull(monitor1);
        assertNotNull(monitor2);

        // cancel remotely
        assertTrue(monitor2.cancel("cancel-token-remotely-running-task-can-be-cancelled", 10, TimeUnit.SECONDS));
        assertTrue(task.isCanceled());
    }

    @Test(expected = IncorrectTokenMaintenanceException.class)
    public void testUnlockWrongToken() {
        mockAuthenticatedAs("user-unlock-wrong-token", 1);
        when(tokenGenerator.generateToken()).thenReturn("token-unlock-wrong-token");

        MaintenanceLock lock = service1.lock();
        lock.unlock("wrong-token");
    }

    @Test(expected = IncorrectTokenMaintenanceException.class)
    public void testUnlockWrongTokenOtherNode() {
        mockAuthenticatedAs("user", 1);
        when(tokenGenerator.generateToken()).thenReturn("token");

        service1.lock();
        MaintenanceLock lock = service2.getLock();

        assertNotNull(lock);
        lock.unlock("wrong-token");
    }

    private DefaultMaintenanceService createMaintenanceService(int nodeIndex) {
        HazelcastInstance hazelcast = cluster.getNode(nodeIndex);

        ClusterService clusterService = mock(ClusterService.class);
        clusterServices[nodeIndex] = clusterService;
        when(clusterService.isClustered()).thenReturn(true);

        String nodeId = hazelcast.getCluster().getLocalMember().getUuid();
        ClusterInformation clusterInfo = mock(ClusterInformation.class);
        ClusterNode node = mock(ClusterNode.class);
        when(node.getId()).thenReturn(nodeId);
        when(clusterInfo.getLocalNode()).thenReturn(node);
        when(clusterService.getInformation()).thenReturn(clusterInfo);
        when(clusterService.getNodeId()).thenReturn(nodeId);


        IExecutorService clusterExecutorService = hazelcast.getExecutorService(HazelcastConstants.EXECUTOR_BITBUCKET_CORE + counter);
        IAtomicReference<ClusterMaintenanceLock> clusterLock = hazelcast.getAtomicReference("cluster-lock-" + counter);
        IAtomicLong isActive = hazelcast.getAtomicLong("task-active-" + counter);

        DefaultMaintenanceService service = new DefaultMaintenanceService(authenticationContext, clusterExecutorService, clusterService,
                databaseManager, eventPublisher, executorService, i18nService, requestManager, scmService,
                tokenGenerator, latestTask, isActive, clusterLock);
        // no delays to speed up testing
        service.setNodeJoinCheckDelayMillis(0, TimeUnit.MILLISECONDS);

        cluster.registerBeansForNode(nodeIndex, clusterExecutorService, i18nService, service);
        return service;
    }

    @SuppressWarnings("ConstantConditions")
    private void doTestStart(LatchingMaintenanceTask task) throws InterruptedException {
        final String sessionId = "ssssssessssionId";
        final String token = "cancel-me";

        RequestContext requestContext = mock(RequestContext.class);
        when(requestContext.getSessionId()).thenReturn(sessionId);
        when(requestManager.getRequestContext()).thenReturn(requestContext);
        when(tokenGenerator.generateToken()).thenReturn(token);

        MaintenanceTaskMonitor taskMonitor = service1.start(task, MaintenanceType.BACKUP);
        assertNotNull(taskMonitor);

        assertEquals(token, taskMonitor.getCancelToken());
        assertTrue(taskMonitor.isOwner(requestContext));
        assertSame(taskMonitor, service1.getRunningTask());

        // verify the status
        MaintenanceStatus status = service1.getStatus();
        assertEquals(MaintenanceTaskState.RUNNING, status.getLatestTask().getState());
        assertEquals(task.getProgress().getMessage(), status.getLatestTask().getProgress().getMessage());

        // let the task finish
        task.release();
        taskMonitor.awaitCompletion();

        // verify that the task is marked as having completed
        assertNull(service1.getRunningTask());

        // verify the status
        status = service1.getStatus();
        assertEquals(MaintenanceTaskState.SUCCESSFUL, status.getLatestTask().getState());
        assertEquals(task.getProgress().getMessage(), status.getLatestTask().getProgress().getMessage());
        assertNull(service1.getRunningTask());
    }

    private void mockAuthenticatedAs(String username, int userId) {
        mockAuthenticatedAs(username, userId, username + "-session-id");
    }

    private void mockAuthenticatedAs(String username, int userId, String sessionId) {
        RequestContext requestContext = mock(RequestContext.class);
        when(requestContext.getSessionId()).thenReturn(sessionId);
        when(requestManager.getRequestContext()).thenReturn(requestContext);
        when(requestContext.getAuthenticationContext()).thenReturn(authenticationContext);
        when(authenticationContext.getCurrentUser()).thenReturn(new SimpleUser(userId, username));
    }

    private void mockStandaloneSetup() {
        reset(clusterServices[0]);
        reset(clusterServices[1]);

        // isClustered and isAvailable return false by default - marking the install as not-clustered
    }

    private static class SimpleUser implements ApplicationUser, Serializable {

        private final int id;
        private final String name;

        private SimpleUser(int id, String name) {
            this.id = id;
            this.name = name;
        }

        @Override
        public <T> T accept(@Nonnull ApplicationUserVisitor<T> visitor) {
            return visitor.visit(this);
        }

        @Nonnull
        @Override
        public String getDisplayName() {
            return name + "-displayname";
        }

        @Override
        public int getId() {
            return id;
        }

        @Nonnull
        @Override
        public UserType getType() {
            return UserType.NORMAL;
        }

        @Override
        public boolean isActive() {
            return true;
        }

        @Nonnull
        @Override
        public String getSlug() {
            return name + "-slug";
        }

        @Override
        public String getEmailAddress() {
            return name + "@example.com";
        }

        @Nonnull
        @Override
        public String getName() {
            return name;
        }
    }

    @ClusterableTask
    private static class ClusterMaintenanceTask extends LatchingMaintenanceTask {
        private ClusterMaintenanceTask(String message, int progress) {
            super(message, progress);
        }
    }

    private static class StandaloneMaintenanceTask extends LatchingMaintenanceTask {
        private StandaloneMaintenanceTask(String message, int progress) {
            super(message, progress);
        }
    }

    private abstract static class LatchingMaintenanceTask implements MaintenanceTask {
        private final CountDownLatch latch = new CountDownLatch(1);

        private volatile boolean canceled;
        private String message;
        private int progress;

        private LatchingMaintenanceTask(String message, int progress) {
            this.message = message;
            this.progress = progress;
        }

        @Override
        public void cancel() {
            canceled = true;
            latch.countDown();
        }

        boolean isCanceled() {
            return canceled;
        }

        @Nonnull
        @Override
        public Progress getProgress() {
            return new ProgressImpl(message, progress);
        }

        public void release() {
            latch.countDown();
        }

        @Override
        public void run() {
            try {
                latch.await();
                throwIfCanceled();
            } catch (InterruptedException e) {
                throwIfCanceled();
                Throwables.propagate(e);
            }
        }

        private void throwIfCanceled() {
            if (canceled) {
                throw new CanceledBackupException(new KeyedMessage("", "", ""));
            }
        }
    }
}
