package com.eactive.testmaster.loadtest.service;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import javax.annotation.PreDestroy;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Scope;
import org.springframework.messaging.simp.SimpMessagingTemplate;
import org.springframework.stereotype.Service;

@Service
@Scope("session")
public class CallableLoadTestService {

    @Autowired
    private SimpMessagingTemplate template;

    private ScheduledExecutorService scheduledExecutorService;
    private List<ScheduledFuture<?>> scheduledFutures;

    public CallableLoadTestService() {
        this.scheduledFutures = new ArrayList<>();
    }

    public void sendLoadTestUpdate(String topic, Object update) {
        template.convertAndSend("/topic/" + topic, update);
    }

    public void startLoadTest(int numberOfUsers, long rampUpTimeMillis, long taskDurationMillis) {
        if (scheduledExecutorService != null && !scheduledExecutorService.isShutdown()) {
            scheduledExecutorService.shutdownNow();
        }
        scheduledExecutorService = Executors.newScheduledThreadPool(numberOfUsers);

        for (int i = 0; i < numberOfUsers; i++) {
            final CallableVirtualUser virtualUser = new CallableVirtualUser(this);

            // Calculate delay for each user to simulate ramp-up
            long delay = i * rampUpTimeMillis;

            // Schedule the task with the calculated delay
            ScheduledFuture<?> future = scheduledExecutorService.schedule(() -> {
                try {
                    // Execute the user task and wait for its completion or interruption
                    String result = virtualUser.call();
                    System.out.println("Task completed with result: " + result);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }, delay, TimeUnit.MILLISECONDS);

            scheduledFutures.add(future);

            // Schedule a stop action for each task after its duration expires
            scheduledExecutorService.schedule(() -> {
                if (!future.isDone()) {
                    future.cancel(true); // Attempt to interrupt the task
                    System.out.println("Task was cancelled after its duration limit.");
                }
            }, delay + taskDurationMillis, TimeUnit.MILLISECONDS);
        }
    }

    public void stopLoadTest() {
        if (scheduledFutures != null) {
            for (ScheduledFuture<?> future : scheduledFutures) {
                if (!future.isDone()) {
                    future.cancel(true); // Attempt to interrupt if running
                }
            }
            scheduledFutures.clear();
        }
        if (scheduledExecutorService != null) {
            scheduledExecutorService.shutdownNow();
        }
    }

    @PreDestroy
    public void cleanup() {
        stopLoadTest(); // Ensure all tasks are stopped and executor service is shut down
    }
}

