package com.agent.management.engine;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;

import java.util.*;

/**
 * DAG 图解析与拓扑排序
 * 解析 graphData JSON，构建邻接表，Kahn 算法分层排序
 */
public class DagResolver {

    private static final ObjectMapper MAPPER = new ObjectMapper();

    @Data
    public static class EdgeInfo {
        private final String source;
        private final String target;
        private final String sourceHandle;

        public EdgeInfo(String source, String target, String sourceHandle) {
            this.source = source;
            this.target = target;
            this.sourceHandle = sourceHandle;
        }
    }

    @Data
    public static class ResolvedDag {
        /** 拓扑层级，每层内的节点无相互依赖 */
        private final List<List<String>> levels = new ArrayList<>();
        /** nodeId → node data (JSON) */
        private final Map<String, JsonNode> nodeDataMap = new HashMap<>();
        /** nodeId → 节点类型 */
        private final Map<String, String> nodeTypeMap = new HashMap<>();
        /** nodeId → 该节点的所有出边 */
        private final Map<String, List<EdgeInfo>> outgoingEdges = new HashMap<>();
        /** nodeId → 该节点的所有入边 */
        private final Map<String, List<EdgeInfo>> incomingEdges = new HashMap<>();
    }

    /**
     * 解析 graphData JSON，返回拓扑排序后的 DAG
     */
    public static ResolvedDag resolve(String graphData) {
        try {
            JsonNode root = MAPPER.readTree(graphData);
            JsonNode nodesArr = root.path("nodes");
            JsonNode edgesArr = root.path("edges");

            ResolvedDag dag = new ResolvedDag();
            Map<String, Integer> inDegree = new HashMap<>();

            // 1. 索引所有节点（防御性规范化 type，应对绕过 Service 的脏数据）
            for (JsonNode node : nodesArr) {
                String id = node.get("id").asText();
                dag.getNodeDataMap().put(id, node.path("data"));
                dag.getNodeTypeMap().put(id, NodeTypeUtils.normalize(node.path("type").asText()));
                dag.getOutgoingEdges().put(id, new ArrayList<>());
                dag.getIncomingEdges().put(id, new ArrayList<>());
                inDegree.put(id, 0);
            }

            // 2. 构建边关系 + 计算入度
            // 用 Set<target> 按源去重，避免同一源节点的多条分支重复计入入度
            Map<String, Set<String>> targetSourceSet = new HashMap<>();
            for (JsonNode edge : edgesArr) {
                String source = edge.get("source").asText();
                String target = edge.get("target").asText();
                String sourceHandle = edge.path("sourceHandle").asText("");
                if (sourceHandle.isEmpty()) sourceHandle = "output";

                EdgeInfo edgeInfo = new EdgeInfo(source, target, sourceHandle);
                dag.getOutgoingEdges().get(source).add(edgeInfo);
                dag.getIncomingEdges().get(target).add(edgeInfo);
                targetSourceSet.computeIfAbsent(target, k -> new HashSet<>()).add(source);
            }
            for (Map.Entry<String, Set<String>> e : targetSourceSet.entrySet()) {
                inDegree.put(e.getKey(), e.getValue().size());
            }

            // 3. Kahn 算法拓扑排序（按层分组）
            Queue<String> queue = new LinkedList<>();
            for (Map.Entry<String, Integer> e : inDegree.entrySet()) {
                if (e.getValue() == 0) queue.add(e.getKey());
            }

            int totalProcessed = 0;
            while (!queue.isEmpty()) {
                List<String> level = new ArrayList<>();
                int size = queue.size();
                for (int i = 0; i < size; i++) {
                    String nodeId = queue.poll();
                    level.add(nodeId);
                    totalProcessed++;

                    for (EdgeInfo edge : dag.getOutgoingEdges().getOrDefault(nodeId, Collections.emptyList())) {
                        int newDeg = inDegree.get(edge.getTarget()) - 1;
                        inDegree.put(edge.getTarget(), newDeg);
                        if (newDeg == 0) queue.add(edge.getTarget());
                    }
                }
                dag.getLevels().add(level);
            }

            // 4. 检测循环
            if (totalProcessed != dag.getNodeDataMap().size()) {
                throw new IllegalStateException("工作流图中存在循环依赖，无法执行");
            }

            return dag;
        } catch (IllegalStateException e) {
            throw e;
        } catch (Exception e) {
            throw new IllegalStateException("解析工作流图数据失败: " + e.getMessage(), e);
        }
    }
}
