package com.cmvr.test.flow.runtime.engine;

import cn.hutool.core.exceptions.ExceptionUtil;
import com.cmvr.test.flow.builder.FlowGraph;
import com.cmvr.test.flow.builder.FlowNodeWrapper;
import com.cmvr.test.flow.builder.TaskKeyBuilder;
import com.cmvr.test.flow.context.TaskContext;
import com.cmvr.test.flow.runtime.engine.support.FlowSchedulerCache;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Set;

@Slf4j
@Component
@RequiredArgsConstructor
public class FlowTaskScheduler {

    @Resource(name = "threadPoolTaskExecutor")
    private final ThreadPoolTaskExecutor executor;

    private final FlowSchedulerCache schedulerCache;

    /**
     * 启动主流程调度（从 start 的所有出边节点开始）
     */
    public void start(FlowGraph graph,
                      TaskContext context,
                      FlowItemExecutor.NodeExecutor executorFunc,
                      Runnable onFinished) {
        String instId = context.getInstId();
        String startNodeId = graph.findStartNodeId();
        List<String> nextNodes = graph.getNextNodes(startNodeId);

        if (nextNodes.isEmpty()) {
            log.warn("start 节点无出边，流程直接结束：instId={}", instId);
            onFinished.run();
            return;
        }

        for (String nextId : nextNodes) {
            scheduleNode(instId, nextId, graph, executorFunc, new ArrayList<>()); // 空 iterations
        }
    }

    /**
     * 启动子图调度（从入度为0的节点开始）
     */
    public void subStart(FlowGraph graph,
                         TaskContext context,
                         FlowItemExecutor.NodeExecutor executorFunc,
                         Runnable onFinished,
                         List<Integer> iterations) {
        String instId = context.getInstId();
        FlowNodeWrapper startNode = graph.getInDegreeZeroNode();
        List<String> nextNodes = graph.getNextNodes(startNode.getNodeId());

        if (nextNodes.isEmpty()) {
            log.warn("子图 start 节点无出边，流程直接结束：instId={}", instId);
            onFinished.run();
            return;
        }

        for (String nextId : nextNodes) {
            scheduleNode(instId, nextId, graph, executorFunc, iterations);
        }
    }

    /**
     * 某节点执行完成后，判断并调度其所有后继节点
     */
    public void markCompleted(String instId,
                              String nodeId,
                              FlowGraph graph,
                              FlowItemExecutor.NodeExecutor executorFunc,
                              List<Integer> iterations) {
        List<String> nextNodes = graph.getNextNodes(nodeId);

        for (String nextId : nextNodes) {
            String key = TaskKeyBuilder.buildKey(instId, nextId, iterations);

            // 使用 Set 来记录已完成的前驱节点
            Set<String> completedPreNodes = schedulerCache.getOrInitCompletedPre(key);
            // 添加当前节点作为已完成的前驱节点
            completedPreNodes.add(nodeId);

            int preCount = graph.getPredecessors().getOrDefault(nextId, Collections.emptyList()).size();

            // 如果当前节点的所有前驱节点都已完成，则调度后继节点
            if (completedPreNodes.size() == preCount) {
                scheduleNode(instId, nextId, graph, executorFunc, iterations);
            }
        }
    }
    /**
     * 提交节点执行任务（除循环只调度一次）
     */
    private void scheduleNode(String instId,
                              String nodeId,
                              FlowGraph graph,
                              FlowItemExecutor.NodeExecutor executorFunc,
                              List<Integer> iterations) {
        String key = TaskKeyBuilder.buildKey(instId, nodeId, iterations);

        // 如果已经调度过，则不再重复调度
        if (schedulerCache.trySchedule(key)) {
            executor.submit(() -> {
                Thread.currentThread().setName("node-thread-" + nodeId);
                try {
                    executorFunc.execute(graph.getNode(nodeId));
                } catch (Exception e) {
                    if (ExceptionUtil.getRootCause(e) instanceof InterruptedException) {
                        return;
                    }
                    log.error("节点执行异常：instId={}, nodeId={}, err={}", instId, nodeId, e.getMessage(), e);
                }
            });
        }
    }

    public void scheduleDirectly(String instId,
                                 String nodeId,
                                 FlowGraph graph,
                                 FlowItemExecutor.NodeExecutor executorFunc,
                                 List<Integer> iterations) {

        String key = TaskKeyBuilder.buildKey(instId, nodeId, iterations);
        if (schedulerCache.trySchedule(key)) {
            executor.submit(() -> {
                try {
                    executorFunc.execute(graph.getNode(nodeId));
                } catch (Exception e) {
                    if (ExceptionUtil.getRootCause(e) instanceof InterruptedException) {
                        return;
                    }
                    log.error("节点执行异常：instId={}, nodeId={}, err={}", instId, nodeId, e.getMessage(), e);
                }
            });
        }
    }

    /**
     * 清理缓存（主图或子图执行结束后调用）
     */
    public void clear(String instId) {
        schedulerCache.clearContextCache(instId);
    }
}
