package com.atlassian.bitbucket.mesh.repository;

import com.atlassian.bitbucket.mesh.network.MeshManager;
import com.atlassian.bitbucket.mesh.partition.Partition;
import com.google.common.collect.ImmutableList;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.io.UncheckedIOException;
import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.function.BiConsumer;

import static java.util.concurrent.TimeUnit.SECONDS;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class ManagedDataStoreCleanupManagerTest {

    @Mock
    private ScheduledExecutorService cleanupExecutor;
    @Mock
    private ManagedDataStore managedDataStore;
    @Mock
    private MeshManager meshManager;
    private List<String> partitions;
    @Captor
    private ArgumentCaptor<Runnable> runnableCaptor;
    @Captor
    private ArgumentCaptor<BiConsumer<Partition, Instant>> streamCaptor;
    private ManagedDataStoreCleanupManager underTest;

    @Before
    public void setup() {
        partitions = ImmutableList.of();
        when(managedDataStore.getAllPartitions()).thenAnswer(invocation -> partitions);
        underTest = new ManagedDataStoreCleanupManager(
                Duration.ofSeconds(60), cleanupExecutor, managedDataStore, meshManager);
    }

    @Test
    public void testExceptionWhileMarkingRepositoriesAsDeletedIsHandled() {
        when(managedDataStore.getAllPartitions()).thenThrow(UncheckedIOException.class);

        underTest.onTopologyUpdatedEvent();

        verify(cleanupExecutor).execute(runnableCaptor.capture());
        runnableCaptor.getValue().run();
    }

    @Test
    public void testRejectedExecutionExceptionsAreHandledForCleanupPartitions() {
        partitions = ImmutableList.of("0000");
        when(managedDataStore.markPartitionDeleted(new Partition("0000"))).thenReturn(true);
        doThrow(RejectedExecutionException.class).when(cleanupExecutor).schedule(any(Runnable.class), anyLong(), any());

        underTest.onTopologyUpdatedEvent();

        verify(cleanupExecutor).execute(runnableCaptor.capture());
        runnableCaptor.getValue().run();

        verify(cleanupExecutor).schedule(any(Runnable.class), anyLong(), any());
    }

    @Test
    public void testRejectedExecutionExceptionsAreHandledForMarkingPartitions() {
        doThrow(RejectedExecutionException.class).when(cleanupExecutor).execute(any());

        underTest.onTopologyUpdatedEvent();

        verify(cleanupExecutor).execute(any());
    }

    @Test
    public void testShutdownShutsDownExecutor() {
        underTest.shutdown();

        verify(cleanupExecutor).shutdown();

        underTest.shutdown();

        verify(cleanupExecutor).shutdown();
    }

    @Test
    public void testTopologyUpdateSchedulesTasksCorrectly() {
        partitions = ImmutableList.of("0000", "0001", "0002", "0004", "000c", "000f");
        mockNetworkPartition("0000");
        mockNetworkPartition("000f");
        mockNetworkPartition("0004");
        Partition partition1 = new Partition("0001");
        Partition partition2 = new Partition("0002");
        Partition partition12 = new Partition("000c");

        //Some mark calls return true to cause scheduling to be kicked off
        when(managedDataStore.markPartitionDeleted(partition1)).thenReturn(true);
        when(managedDataStore.markPartitionDeleted(partition2)).thenReturn(true);

        underTest.onTopologyUpdatedEvent();

        verify(cleanupExecutor).execute(runnableCaptor.capture());

        //Run the "mark" task
        runnableCaptor.getValue().run();

        verify(meshManager).getPartition("0000");
        verify(meshManager).getPartition("0001");
        verify(meshManager).getPartition("0002");
        verify(meshManager).getPartition("0004");
        verify(meshManager).getPartition("000c");
        verify(meshManager).getPartition("000f");

        verify(managedDataStore).getAllPartitions();
        verify(managedDataStore).markPartitionDeleted(partition2);
        verify(managedDataStore).markPartitionDeleted(partition1);
        verify(managedDataStore).markPartitionDeleted(partition12);
        verifyNoMoreInteractions(managedDataStore, meshManager);

        verify(cleanupExecutor).schedule(runnableCaptor.capture(), eq(60L), eq(SECONDS));

        //Run the "cleanup" task
        runnableCaptor.getValue().run();
        Instant now = Instant.now();

        verify(managedDataStore).streamDeletedPartitions(streamCaptor.capture());
        BiConsumer<Partition, Instant> consumer = streamCaptor.getValue();

        //Mark partition 0001 as preset again, should not get deleted
        //Partition 0002 has been marked less than 1 minute ago and should be skipped
        //Partition 000c should be deleted as it was marked more than 1 minute ago
        mockNetworkPartition(partition1.getHex());

        consumer.accept(partition1, now.minusSeconds(61));
        consumer.accept(partition2, now.minusSeconds(30));
        consumer.accept(partition12, now.minusSeconds(61));

        verify(managedDataStore).clearPartitionDeleted(partition1);
        verify(managedDataStore).deletePartitionFiles(partition12);
        verifyNoMoreInteractions(managedDataStore);
    }

    @Test
    public void testTopologyUpdateWithNoFilesMarkedDoesntScheduleCleanup() {
        Partition partition = new Partition("0000");
        partitions = ImmutableList.of(partition.getHex());

        //Some mark calls return true to cause scheduling to be kicked off
        when(managedDataStore.markPartitionDeleted(partition)).thenReturn(false);

        underTest.onTopologyUpdatedEvent();

        verify(cleanupExecutor).execute(runnableCaptor.capture());

        //Run the "mark" task
        runnableCaptor.getValue().run();

        verify(meshManager).getPartition("0000");

        verify(managedDataStore).getAllPartitions();
        verify(managedDataStore).markPartitionDeleted(partition);
        verifyNoMoreInteractions(managedDataStore, meshManager);

        verify(cleanupExecutor, never()).schedule(any(Runnable.class), anyLong(), any());

        verifyNoMoreInteractions(cleanupExecutor);
    }

    private void mockNetworkPartition(String id) {
        when(meshManager.getPartition(id)).thenReturn(
                new com.atlassian.bitbucket.mesh.network.Partition.Builder(id).build());
    }
}