package com.eactive.testmaster.loadtest.service;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;
import javax.annotation.PreDestroy;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.stereotype.Component;

@Component
public class LoadTestMonitorMap {

    private final Map<String, List<ScheduledFuture<?>>> futures = new ConcurrentHashMap<>();
    private final Map<String, List<VirtualUser>> virtualUsers = new ConcurrentHashMap<>();

    private final Map<String, ThreadPoolTaskScheduler> schedulers = new ConcurrentHashMap<>();

    public void put(String testId, List<ScheduledFuture<?>> futures, List<VirtualUser> users, ThreadPoolTaskScheduler scheduler) {
        this.futures.put(testId, futures);
        virtualUsers.put(testId, users);
        schedulers.put(testId, scheduler);
    }

    public void remove(String testId) {
        futures.remove(testId);
        virtualUsers.remove(testId);
        schedulers.remove(testId);
    }

    public ThreadPoolTaskScheduler getScheduler(String testId) {
        return schedulers.get(testId);
    }

    public List<ScheduledFuture<?>> getFutures(String testId) {
        return futures.get(testId);
    }

    public List<VirtualUser> getVirtualUsers(String testId) {
        return virtualUsers.get(testId);
    }

    @PreDestroy
    public void cleanup() {
        // Perform cleanup logic here
        //get all the keys and stop them all
        futures.keySet().forEach(key -> {
            List<ScheduledFuture<?>> value = futures.get(key);
            value.forEach(future -> {
                future.cancel(true);
            });
        });

        virtualUsers.keySet().forEach(key -> {
            List<VirtualUser> value = virtualUsers.get(key);
            value.forEach(VirtualUser::stop);
        });

        schedulers.keySet().forEach(key -> {
            ThreadPoolTaskScheduler value = schedulers.get(key);
            value.shutdown();
        });

    }
}
