package com.agent.management.engine.executor;

import com.agent.management.engine.*;
import com.agent.management.service.AiModelService;
import com.fasterxml.jackson.databind.JsonNode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.stream.Collectors;

/**
 * 条件分支节点执行器
 * 使用 LLM 根据上下文数据判断走哪个分支，返回选中的 sourceHandle
 */
@Slf4j
@Component
public class ConditionExecutor implements NodeExecutor {

    private final AiModelService aiModelService;
    private final ChatClient.Builder fallbackBuilder;

    public ConditionExecutor(ChatClient.Builder chatClientBuilder, AiModelService aiModelService) {
        this.fallbackBuilder = chatClientBuilder;
        this.aiModelService = aiModelService;
    }

    @Override
    public String getType() {
        return "condition";
    }

    @Override
    public NodeExecutionResult execute(String nodeId, JsonNode data, WorkflowContext context, NodeWorkspace workspace) {
        JsonNode conditions = data.path("conditions");
        if (!conditions.isArray() || conditions.isEmpty()) {
            // 没有条件定义，走默认分支（路由靠 sourceHandle，selectedBranch 放入 data 方便下游读取）
            Map<String, Object> dataMap = new java.util.LinkedHashMap<>();
            dataMap.put("selectedBranch", "branch-0");
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.success("条件路由完成：branch-0（无分支定义，走默认）", dataMap).toMap(),
                    "branch-0");
        }

        // 构建条件列表文本
        StringBuilder conditionText = new StringBuilder();
        int branchCount = 0;
        for (JsonNode cond : conditions) {
            String type = cond.path("type").asText("IF");
            String expression = cond.path("expression").asText("");
            conditionText.append(branchCount).append(" (").append(type).append("): ").append(expression).append("\n");
            branchCount++;
        }
        // 隐含 ELSE 分支
        conditionText.append(branchCount).append(" (ELSE): 以上条件都不满足\n");

        // 构建上下文摘要
        String contextSummary = workspace.getVariables().entrySet().stream()
                .filter(e -> e.getValue() != null)
                .map(e -> e.getKey() + ": " + truncate(e.getValue().toString(), 200))
                .collect(Collectors.joining("\n"));

        String prompt = "你是条件判断助手。根据提供的数据判断哪个条件成立，只回复条件编号数字。\n\n"
                + "数据:\n" + (contextSummary.isEmpty() ? "(无)" : contextSummary) + "\n\n"
                + "条件:\n" + conditionText + "\n"
                + "请只回复编号数字(0-" + branchCount + ")，不要回复其他内容。";

        log.info("[Condition] 节点 {} 开始判断", nodeId);

        try {
            String response = aiModelService.getChatClientOrFallback(fallbackBuilder)
                    .prompt()
                    .user(prompt)
                    .call()
                    .content();

            int selectedBranch = parseBranchIndex(response, branchCount);
            String sourceHandle = "branch-" + selectedBranch;
            log.info("[Condition] 节点 {} 选择分支: {} (response={})", nodeId, sourceHandle, response);

            Map<String, Object> dataMap = new java.util.LinkedHashMap<>();
            dataMap.put("selectedBranch", sourceHandle);
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.success("条件路由完成：" + sourceHandle, dataMap).toMap(),
                    sourceHandle);

        } catch (Exception e) {
            log.error("[Condition] 节点 {} 判断失败: {}", nodeId, e.getMessage());
            // 降级：走第一个分支，但标记为失败 envelope
            Map<String, Object> dataMap = new java.util.LinkedHashMap<>();
            dataMap.put("selectedBranch", "branch-0");
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.of(NodeOutputEnvelope.STATUS_FAILED,
                            "条件判断失败，降级走 branch-0：" + e.getMessage(), dataMap).toMap(),
                    "branch-0");
        }
    }

    /**
     * 解析 LLM 返回的分支编号
     */
    private int parseBranchIndex(String response, int maxBranch) {
        if (response == null) return 0;
        // 提取数字
        String trimmed = response.trim().replaceAll("[^0-9]", "");
        if (trimmed.isEmpty()) return 0;
        try {
            int idx = Integer.parseInt(trimmed);
            return Math.min(idx, maxBranch);
        } catch (NumberFormatException e) {
            return 0;
        }
    }

    private String truncate(String text, int maxLen) {
        if (text == null) return "";
        return text.length() <= maxLen ? text : text.substring(0, maxLen) + "...";
    }
}
