package com.atlassian.bitbucket.mesh.repair;

import com.atlassian.bitbucket.mesh.client.MeshClient;
import com.atlassian.bitbucket.mesh.config.Backoff;
import com.atlassian.bitbucket.mesh.config.RepairConfig;
import com.atlassian.bitbucket.mesh.execution.GrpcExecutionManager;
import com.atlassian.bitbucket.mesh.git.exception.RepositoryRepairFailedException;
import com.atlassian.bitbucket.mesh.junit.LogListener;
import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.network.MeshNode;
import com.atlassian.bitbucket.mesh.network.Partition;
import com.atlassian.bitbucket.mesh.network.PartitionReplica;
import com.atlassian.bitbucket.mesh.repository.*;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.MoreExecutors;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;
import org.springframework.scheduling.TaskScheduler;

import java.time.Duration;
import java.time.Instant;
import java.util.List;

import static com.atlassian.bitbucket.mesh.repository.ReplicaState.*;
import static java.util.Collections.singleton;
import static org.junit.Assert.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class DefaultRepairManagerTest {

    private static final String PARTITION = "0023";
    private static final String REPOSITORY_ID = "p/" + PARTITION + "/h/7236ce673d92c3960ae62/r/2623";
    private static final long VERSION = 923L;

    @Rule
    public final LogListener logListener = new LogListener(DefaultRepairManager.class);

    @Mock
    private Backoff backoff;
    @Mock
    private RepairConfig config;
    @Mock
    private MeshClient client;
    @Mock
    private MeshNode localNode;
    @Mock
    private MeshManager meshManager;
    @Mock
    private MeshNode node1;
    @Mock
    private MeshNode node2;
    @Mock
    private ReplicaStateRegistry registry;
    private DefaultRepairManager repairManager;
    @Mock
    private TaskScheduler taskScheduler;

    @Before
    public void setUp() {
        when(config.getBackoff()).thenReturn(backoff);
        when(config.getMaxAttempts()).thenReturn(3);
        when(config.getTimeout()).thenReturn(Duration.ofSeconds(1L));
        when(registry.getState(REPOSITORY_ID)).thenReturn(INCONSISTENT);

        Backoff.Execution backoffExecution = mock(Backoff.Execution.class);
        when(backoff.start()).thenReturn(backoffExecution);
        when(backoffExecution.nextDelay()).thenReturn(Duration.ZERO);

        Partition partition = mock(Partition.class);
        List<PartitionReplica> replicas = ImmutableList.of(
                mockReplica(localNode), mockReplica(node1), mockReplica(node2));
        when(partition.getReplicas()).thenReturn(replicas);
        when(meshManager.getPartition(PARTITION)).thenReturn(partition);
        when(localNode.isLocal()).thenReturn(true);
        when(node1.isRemote()).thenReturn(true);
        when(node2.isRemote()).thenReturn(true);

        when(taskScheduler.schedule(any(), any(Instant.class))).thenAnswer(invocation -> {
            Runnable task = invocation.getArgument(0);
            task.run();
            return null;
        });

        repairManager = new DefaultRepairManager(config, new GrpcExecutionManager(),
                MoreExecutors.newDirectExecutorService(), client, meshManager, registry, taskScheduler);
    }

    @After
    public void tearDown() {
        repairManager.destroy();
    }

    @Test
    public void testDirectedRepairForceRepairs() {
        //Mark the repository as consistent according to the registry
        lenient().when(registry.getState(REPOSITORY_ID)).thenReturn(CONSISTENT);
        //Mock a successful repair
        when(client.repairMe(eq(node1), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), CONSISTENT, CONSISTENT, VERSION));

        assertEquals(CONSISTENT, repairManager.repair(REPOSITORY_ID, singleton(node1), localNode));
    }

    @Test
    public void testDirectedRepairReturnsInconsistentWhenInconsistentAfterRepair() {
        //Mark the repository as consistent according to the registry
        lenient().when(registry.getState(REPOSITORY_ID)).thenReturn(CONSISTENT);
        //Mock a failed repair
        when(client.repairMe(eq(node1), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), CONSISTENT, INCONSISTENT, VERSION));

        assertEquals(INCONSISTENT, repairManager.repair(REPOSITORY_ID, singleton(node1), localNode));
    }

    @Test
    public void testDirectedRepairReturnsRepairingWhenRepairIsAlreadyRunning() {
        lenient().when(registry.getState(REPOSITORY_ID)).thenReturn(INCONSISTENT);
        when(client.repairMe(any(), any())).thenThrow(RepositoryRepairAlreadyRunningException.class);

        assertEquals(REPAIRING, repairManager.repair(REPOSITORY_ID, singleton(node1), localNode));
    }

    @Test
    public void testDirectedRepairWhenAllRepairSlotsAreTaken() {
        //Mark the repository as consistent according to the registry
        lenient().when(registry.getState(REPOSITORY_ID)).thenReturn(CONSISTENT);
        //Mock a failed repair
        RepositoryRepairSlotsExhaustedException expected = new RepositoryRepairSlotsExhaustedException(
                "Simulated all repair slots busy");
        when(client.repairMe(eq(node1), any()))
                .thenThrow(expected);

        RepositoryRepairFailedException actualException = assertThrows(RepositoryRepairFailedException.class,
                () -> repairManager.repair(REPOSITORY_ID, singleton(node1), localNode));
        assertNotNull(actualException);
        assertEquals(expected, actualException.getCause());

        //Verify that no retry has been scheduled
        verifyNoInteractions(taskScheduler);
    }

    @Test
    public void testDirectedRepairThrowsWhenAllReplicasFail() {
        //Mark the repository as consistent according to the registry
        lenient().when(registry.getState(REPOSITORY_ID)).thenReturn(CONSISTENT);
        //Mock a failed repair
        RepositoryRepairFailedException expected = new RepositoryRepairFailedException("Simulated repair failure");
        when(client.repairMe(eq(node1), any()))
                .thenThrow(expected);

        RepositoryRepairFailedException actualException = assertThrows(RepositoryRepairFailedException.class,
                () -> repairManager.repair(REPOSITORY_ID, singleton(node1), localNode));
        assertNotNull(actualException);
        assertEquals(expected, actualException.getCause());

        //Verify that no retry has been scheduled
        verifyNoInteractions(taskScheduler);
    }

    @Test
    public void testOnReplicaStateObservedClearsAbandonedButDoesNotRepairIfConsistent() {
        //Train inconsistent repairs to setup abandonment
        when(client.repairMe(any(), any())).thenReturn(
                new RepairSummary(Instant.now(), CONSISTENT, INCONSISTENT, VERSION));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));
        assertTrue("Repairs should be abandoned", repairManager.isAbandoned(REPOSITORY_ID));

        reset(client, registry, taskScheduler);

        //Now simulate the replica (somehow) being observed as consistent
        repairManager.onReplicaStateObserved(new ReplicaStateObservedEvent(REPOSITORY_ID,
                new ReplicaStateObservation(CONSISTENT, System.currentTimeMillis(), VERSION)));
        assertFalse("Abandoned flag should be cleared", repairManager.isAbandoned(REPOSITORY_ID));

        //Verify no repair was attempted
        verifyNoInteractions(client, registry, taskScheduler);
    }

    @Test
    public void testOnReplicaStateObservedIsIgnoredIfNotAbandoned() {
        repairManager.onReplicaStateObserved(new ReplicaStateObservedEvent(REPOSITORY_ID,
                new ReplicaStateObservation(INCONSISTENT, System.currentTimeMillis(), VERSION)));

        verifyNoInteractions(client, registry, taskScheduler);
    }

    @Test
    public void testOnRepositoryDeletedClearsAbandoned() {
        //Train inconsistent repairs to setup abandonment
        when(client.repairMe(any(), any())).thenReturn(
                new RepairSummary(Instant.now(), CONSISTENT, INCONSISTENT, VERSION));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));
        assertTrue("Repairs should be abandoned", repairManager.isAbandoned(REPOSITORY_ID));

        reset(client, registry, taskScheduler);

        //Now simulate the repository being deleted
        Repository repository = mock(Repository.class);
        when(repository.getId()).thenReturn(REPOSITORY_ID);

        repairManager.onRepositoryDeleted(new RepositoryDeletedEvent(repository));
        assertFalse("Abandoned flag should be cleared", repairManager.isAbandoned(REPOSITORY_ID));

        //Verify we don't attempt to trigger repair for a deleted repository
        verifyNoInteractions(client, registry, taskScheduler);
    }

    @Test
    public void testRepairIsAbandonedIfInconsistentAfterMaxAttempts() {
        when(client.repairMe(any(), any())).thenReturn(
                new RepairSummary(Instant.now(), CONSISTENT, INCONSISTENT, VERSION));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));
        assertTrue("Repairs should be abandoned", repairManager.isAbandoned(REPOSITORY_ID));

        verify(client, times(3)).repairMe(any(), any());
        //The repair should be retried immediately (without a delay) because repair was initiated, did not
        //result in an up-to-date replica, and the repair source _is_ consistent. Under these circumstances,
        //there's no reason to delay the retry.
        verifyNoInteractions(taskScheduler);
    }

    @Test
    public void testRepairIsSkippedIfAlreadyRunning() {
        when(registry.getState(REPOSITORY_ID)).thenReturn(INCONSISTENT);
        when(client.repairMe(any(), any())).thenThrow(RepositoryRepairAlreadyRunningException.class);

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));

        //Repair should not be retried, nor scheduled to be retried later
        verify(client).repairMe(any(), any());
        verifyNoInteractions(taskScheduler);
    }

    @Test
    public void testRepairIsSkippedIfConsistentAccordingToRegistry() {
        when(registry.getState(REPOSITORY_ID)).thenReturn(CONSISTENT);

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));

        verifyNoInteractions(client);
    }

    @Test
    public void testRepairIsRescheduledIfInconsistentAfterRepair() {
        //Simulate an unsuccessful repair followed by a successful repair
        when(client.repairMe(any(), any())).thenReturn(
                new RepairSummary(Instant.now(), CONSISTENT, INCONSISTENT, VERSION),
                new RepairSummary(Instant.now(), CONSISTENT, CONSISTENT, VERSION));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));

        verify(client, times(2)).repairMe(any(), any());
        //The repair should be retried immediately (without a delay) because repair was initiated, did not
        //result in an up-to-date replica, and the repair source _is_ consistent. Under these circumstances,
        //there's no reason to delay the retry.
        verifyNoInteractions(taskScheduler);
    }

    @Test
    public void testRepairOtherNodeFailure() {
        //Mark the repository as consistent according to the registry
        when(registry.getState(REPOSITORY_ID)).thenReturn(CONSISTENT);
        //Mock a failed repair
        when(client.repair(eq(node1), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), CONSISTENT, INCONSISTENT, VERSION));

        assertEquals(INCONSISTENT, repairManager.repair(REPOSITORY_ID, singleton(localNode), node1));
    }

    @Test
    public void testRepairOtherNodeSuccess() {
        //Mark the repository as consistent according to the registry
        when(registry.getState(REPOSITORY_ID)).thenReturn(CONSISTENT);
        //Mock a successful repair
        when(client.repair(eq(node1), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), CONSISTENT, CONSISTENT, VERSION));

        assertEquals(CONSISTENT, repairManager.repair(REPOSITORY_ID, singleton(localNode), node1));
    }

    @Test
    public void testRepairOtherNodeThrowsIfLocalNodeInconsistent() {
        //Mark the repository as inconsistent according to the registry
        when(registry.getState(REPOSITORY_ID)).thenReturn(INCONSISTENT);

        //If the local node is not consistent, no repair should be attempted and UNKNOWN should be returned
        assertEquals(ReplicaState.UNKNOWN, repairManager.repair(REPOSITORY_ID, singleton(localNode), node1));
        verifyNoInteractions(client);
    }

    @Test
    public void testRepairTriesNextReplicaIfReplicaIsNotConsistent() {
        //Node 1 reports being inconsistent
        when(client.repairMe(same(node1), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), INCONSISTENT, INCONSISTENT, VERSION));
        //But node 2 reports is consistent
        when(client.repairMe(same(node2), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), CONSISTENT, CONSISTENT, VERSION));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));

        verify(client).repairMe(same(node1), any());
        verify(client).repairMe(same(node2), any());
        verifyNoMoreInteractions(client);
    }

    @Test
    public void testRepairTriesNextReplicaIfReplicaThrows() {
        //Simulate failure when attempting to repair from node 1
        when(client.repairMe(same(node1), any()))
                .thenThrow(new RuntimeException("Simulating MeshClient failure"));
        //But node 2 reports is consistent
        when(client.repairMe(same(node2), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), CONSISTENT, CONSISTENT, VERSION));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));

        verify(client).repairMe(same(node1), any());
        verify(client).repairMe(same(node2), any());
        verifyNoMoreInteractions(client);
    }

    @Test
    public void testRepairIsAbandonedIfAllReplicasFail() {
        //Expect 2 repair attempts (and both should fail)
        //Simulate failure when attempting to repair from any node
        when(client.repairMe(any(), any()))
                .thenThrow(new RuntimeException("Simulating MeshClient failure"));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));
        assertTrue("Repairs should be abandoned", repairManager.isAbandoned(REPOSITORY_ID));

        //Repair should be attempted 3 times, with increasing delays.
        //Each attempt should attempt to repair from each remote node.
        verify(client, times(3)).repairMe(same(node1), any());
        verify(client, times(3)).repairMe(same(node2), any());
        verify(taskScheduler, times(2)).schedule(any(), any(Instant.class));
        verifyNoMoreInteractions(client);

        logListener.assertWarn("Abandoning repair");
    }

    @Test
    public void testRepairIsRestartedOnActivityIfAbandoned() {
        //Train exceptions to set up the failures necessary to abandon repairs
        when(client.repairMe(any(), any()))
                .thenThrow(new RuntimeException("Simulating MeshClient failure"));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));
        assertTrue("Repairs should be abandoned", repairManager.isAbandoned(REPOSITORY_ID));

        logListener.assertWarn("Abandoning repair");

        //Purge the invocations and train a successful repair
        reset(client, taskScheduler);
        when(client.repairMe(any(), any()))
                .thenReturn(new RepairSummary(Instant.now(), CONSISTENT, CONSISTENT, VERSION));

        //Now simulate repository activity in the form of a new observation. This mimics a new failed
        //transaction on the inconsistent replica
        repairManager.onReplicaStateObserved(new ReplicaStateObservedEvent(REPOSITORY_ID,
                new ReplicaStateObservation(INCONSISTENT, System.currentTimeMillis(), VERSION)));

        verify(client).repairMe(same(node1), any());
        verifyNoMoreInteractions(client);
        verifyNoInteractions(taskScheduler);
    }

    @Test
    public void testRepairThrowsIfSourceAndTargetLocal() {
        assertThrows(IllegalArgumentException.class,
                () -> repairManager.repair(REPOSITORY_ID, singleton(localNode), localNode));
    }

    @Test
    public void testRepairThrowsIfSourceAndTargetRemote() {
        assertThrows(IllegalArgumentException.class,
                () -> repairManager.repair(REPOSITORY_ID, singleton(node1), node2));
    }

    @Test
    public void testStateUpdatedEventTriggersRepair() {
        //Simulate a successful repair
        when(client.repairMe(any(), any())).thenAnswer(
                invocation -> new RepairSummary(Instant.now(), CONSISTENT, CONSISTENT, VERSION));

        repairManager.onReplicaStateUpdated(new ReplicaStateUpdatedEvent(REPOSITORY_ID, INCONSISTENT));

        verify(client).repairMe(any(), any());
    }

    private static PartitionReplica mockReplica(MeshNode node) {
        PartitionReplica replica = mock(PartitionReplica.class);
        when(replica.getNode()).thenReturn(node);
        return replica;
    }
}