package dst.ass3.elastic;

import dst.ass3.elastic.impl.ElasticityFactory;
import dst.ass3.messaging.IWorkloadMonitor;
import dst.ass3.messaging.RequestType;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;

@RunWith(MockitoJUnitRunner.class)
public class ElasticityControllerTest {
    private static final String WORKER_IMAGE = "dst/ass3-worker";

    IElasticityFactory factory;

    @Mock
    IContainerService containerService;
    @Mock
    IWorkloadMonitor workloadMonitor;

    IElasticityController elasticityController;
    Map<RequestType, Double> processingTimes = new HashMap<>();
    Map<RequestType, Long> workerCount = new HashMap<>();
    Map<RequestType, Long> requestCount = new HashMap<>();
    List<ContainerInfo> containers = new ArrayList<>();

    @Before
    public void setUp() throws Exception {
        factory = new ElasticityFactory();
        elasticityController = factory.createElasticityController(containerService, workloadMonitor);

        processingTimes.clear();
        processingTimes.put(RequestType.DOCUMENT, 5000.0);
        processingTimes.put(RequestType.VIDEO, 10000.0);
        processingTimes.put(RequestType.QUIZ, 2000.0);
        when(workloadMonitor.getAverageProcessingTime()).thenReturn(processingTimes);

        workerCount.clear();
        workerCount.put(RequestType.DOCUMENT, 95L);
        workerCount.put(RequestType.VIDEO, 87L);
        workerCount.put(RequestType.QUIZ, 61L);
        when(workloadMonitor.getWorkerCount()).thenReturn(workerCount);

        requestCount.clear();
        requestCount.put(RequestType.DOCUMENT, 600L);
        requestCount.put(RequestType.VIDEO, 1000L);
        requestCount.put(RequestType.QUIZ, 1005L);
        when(workloadMonitor.getRequestCount()).thenReturn(requestCount);

        containers.clear();
        for (int i = 0; i < 95; i++) {
            containers.add(containerInfo("documentId" + i, WORKER_IMAGE, RequestType.DOCUMENT, true));
        }
        for (int i = 0; i < 87; i++) {
            containers.add(containerInfo("video" + i, WORKER_IMAGE, RequestType.VIDEO, true));
        }
        for (int i = 0; i < 61; i++) {
            containers.add(containerInfo("quiz" + i, WORKER_IMAGE, RequestType.QUIZ, true));
        }
        when(containerService.listContainers()).thenReturn(containers);
    }

    @After
    public void tearDown() {
        verify(workloadMonitor, atLeast(1)).getWorkerCount();
        verify(workloadMonitor, atLeast(1)).getRequestCount();
        verify(workloadMonitor, atLeast(1)).getAverageProcessingTime();
    }

    @Test
    public void notEnoughWorkers_scaleUp() throws Exception {
        // remove 10 document workers and 10 quiz workers
        List<ContainerInfo> containersToRemove = containers.stream()
                .filter(c -> c.getContainerId().startsWith("documentId7") || c.getContainerId().startsWith("quiz1"))
                .collect(Collectors.toList());
        containers.removeAll(containersToRemove);
        workerCount.put(RequestType.DOCUMENT, 85L);
        workerCount.put(RequestType.QUIZ, 51L);

        elasticityController.adjustWorkers();

        verify(containerService, never()).stopContainer((String) any(String.class));
        verify(containerService, times(15)).startWorker(RequestType.DOCUMENT);
        verify(containerService, times(16)).startWorker(RequestType.QUIZ);
        verify(containerService, never()).startWorker(RequestType.VIDEO);
        verify(containerService, never()).listContainers();
    }

    @Test
    public void tooManyWorkers_scaleDown() throws Exception {
        // add 20 more, some should be stopped
        for (int i = 0; i < 20; i++) {
            containers.add(containerInfo("quiz1" + i, WORKER_IMAGE, RequestType.QUIZ, true));
        }
        workerCount.put(RequestType.QUIZ, 81L);

        elasticityController.adjustWorkers();

        verify(containerService, times(14)).stopContainer(contains("quiz"));
        verify(containerService, never()).stopContainer(contains("video"));
        verify(containerService, never()).stopContainer(contains("document"));
        verify(containerService, never()).startWorker((RequestType) any(RequestType.class));
        verify(containerService, times(1)).listContainers();
    }

    @Test
    public void justEnoughWorkers_doNotScale() throws Exception {
        elasticityController.adjustWorkers();
        verify(containerService, never()).startWorker((RequestType) any(RequestType.class));
        verify(containerService, never()).stopContainer((String) any());
        verify(containerService, never()).listContainers();
    }

    private ContainerInfo containerInfo(String id, String image, RequestType workerType, boolean running) {
        ContainerInfo info = new ContainerInfo();
        info.setContainerId(id);
        info.setImage(image);
        info.setWorkerType(workerType);
        info.setRunning(running);
        return info;
    }

}