package com.dora.scheduler;

import com.dora.entity.PointsConfig;
import com.dora.mapper.PointsConfigMapper;
import com.dora.service.AiTaskService;
import com.dora.service.QueueService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.List;
import java.util.stream.Collectors;
import com.dora.entity.AiTask;
import com.dora.mapper.AiTaskMapper;

@Component
@Slf4j
@RequiredArgsConstructor
public class TaskScheduler {

    private final QueueService queueService;
    private final AiTaskService aiTaskService; // Changed from AsyncTaskExecutor
    private final PointsConfigMapper pointsConfigMapper;
    private final AiTaskMapper aiTaskMapper; // 新增

    @Value("${ai.queue.max-concurrent:10}")
    private int maxConcurrent;

    @Value("${ai.task.timeout-minutes:10}")
    private int taskTimeoutMinutes;


    @Scheduled(fixedRateString = "${ai.queue.scan-interval:5000}")
    public void dispatchTasks() {
        log.trace("执行任务调度器...");

        List<String> activeModels = getActiveModels();
        if (activeModels.isEmpty()) {
            return;
        }

        for (String modelName : activeModels) {
            try {
                dispatchModelTasks(modelName, maxConcurrent);
            } catch (Exception e) {
                log.error("调度模型 '{}' 时发生未知异常。", modelName, e);
            }
        }
    }

    @Scheduled(fixedRateString = "${ai.task.timeout-scan-interval:60000}") // 每分钟检查一次
    public void checkTasksTimeout() {
        log.trace("执行任务超时检查器...");
        LocalDateTime timeoutThreshold = LocalDateTime.now().minusMinutes(taskTimeoutMinutes);

        // 查找所有长时间处于 "processing" 状态的任务
        List<AiTask> stuckTasks = aiTaskMapper.findStuckTasks(timeoutThreshold);

        if (stuckTasks.isEmpty()) {
            return;
        }

        log.warn("发现 {} 个超时任务，将把它们标记为失败。", stuckTasks.size());

        for (AiTask task : stuckTasks) {
            log.warn("任务 {} (模型: {}) 已超时。最后更新时间: {}", task.getTaskNo(), task.getModelName(), task.getUpdateTime());
            // 将任务标记为失败，本次不消耗积分
            // 注意：这里需要并发安全，因为任务可能在检查和标记之间恰好完成了
            aiTaskService.markTaskAsFailed(task.getTaskNo(), "任务处理超时，已自动取消，本次不消耗积分。");
        }
    }

    private void dispatchModelTasks(String modelName, int maxConcurrent) {
        long queueLength = queueService.getQueueLength(modelName);
        if (queueLength == 0) {
            return;
        }

        int processingCount = queueService.getProcessingCount(modelName);
        if (processingCount >= maxConcurrent) {
            return; // 处理槽位已满
        }
        
        log.debug("调度模型: '{}'。正在处理: {}, 队列中: {}, 容量: {}", modelName, processingCount, queueLength, maxConcurrent);

        int tasksToDispatch = (int) Math.min(maxConcurrent - processingCount, queueLength);
        
        for (int i = 0; i < tasksToDispatch; i++) {
            // 改进的并发控制：只有一个线程能成功增加计数并获取任务
            if (queueService.incrementProcessingCount(modelName) > maxConcurrent) {
                queueService.decrementProcessingCount(modelName);
                break; // 其他线程已占满槽位
            }

            String taskNo = queueService.dequeue(modelName);
            if (taskNo == null) {
                queueService.decrementProcessingCount(modelName);
                break; // 队列已空
            }

            log.info("准备派发任务: {}", taskNo);
            try {
                // 现在调用AiTaskService，而不是直接调用Executor
                aiTaskService.processTask(taskNo);
            } catch (Exception e) {
                log.error("提交任务 {} 到服务层失败。", taskNo, e);
                // 关键：提交失败，必须立即释放占位并标记任务失败
                queueService.decrementProcessingCount(modelName);
                aiTaskService.markTaskAsFailed(taskNo, "任务派发失败，无法提交处理。");
            }
        }
    }

    private List<String> getActiveModels() {
        // 此处可以增加缓存以避免频繁查询数据库
        return pointsConfigMapper.findAll().stream()
                .filter(config -> config.getIsEnabled() == 1)
                .map(PointsConfig::getModelName)
                .collect(Collectors.toList());
    }
}
