package com.zksy.data.utils;

import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.support.CronTrigger;
import org.springframework.stereotype.Component;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledFuture;

@Component
public class TaskManager {
    private final ThreadPoolTaskScheduler scheduler;
    private final ConcurrentHashMap<String, ScheduledFuture<?>> taskMap = new ConcurrentHashMap<>();

    public TaskManager(ThreadPoolTaskScheduler scheduler) {
        this.scheduler = scheduler;
    }

    public void addTask(String taskId, Runnable task, String cronExpression) {
        // 检查任务ID是否已经存在
        if (taskMap.containsKey(taskId)) {
            // 如果任务ID已存在，可以选择抛出异常或返回错误信息
            throw new IllegalArgumentException("Task with ID " + taskId + " already exists.");
            // 或者可以选择先移除旧的任务
            // removeTask(taskId);
        }
        ScheduledFuture<?> future = scheduler.schedule(task, new CronTrigger(cronExpression));
        taskMap.put(taskId, future);
    }

    public void removeTask(String taskId) {
        ScheduledFuture<?> future = taskMap.remove(taskId);
        if (future != null) {
            future.cancel(true);
        }
    }

    // 查看所有任务的状态
    public ConcurrentHashMap<String, ScheduledFuture<?>> getTaskMap() {
        return taskMap;
    }

    // 查看单个任务的状态
    public ScheduledFuture<?> getTaskStatus(String taskId) {
        return taskMap.get(taskId);
    }

    // 停止所有定时任务
    public void stopAllTasks() {
        for (ScheduledFuture<?> future : taskMap.values()) {
            if (future != null) {
                future.cancel(true);
            }
        }
        taskMap.clear(); // 清空任务映射，如果你希望在取消后移除任务引用
    }

}