package com.agent.management.engine.executor;

import com.agent.management.engine.*;
import com.agent.management.config.WorkflowProperties;
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.LinkedHashMap;
import java.util.Map;

/**
 * LLM 节点执行器
 * 渲染 systemPrompt/userPrompt 中的 {{变量名}}，调用 ChatClient，结果写入上下文
 * 支持 modelId 字段指定模型，无 modelId 时使用默认模型
 */
@Slf4j
@Component
public class LlmExecutor implements NodeExecutor {

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

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

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

    @Override
    public NodeExecutionResult execute(String nodeId, JsonNode data, WorkflowContext context, NodeWorkspace workspace) {
        // 系统提示词：渲染用户自定义 prompt，并仅注入当前节点明确定义的输入变量
        String systemPrompt = ContextPromptHelper.merge(
                TemplateRenderer.render(data.path("systemPrompt").asText(""), workspace.getVariables()),
                data, workspace);
        String userPrompt = TemplateRenderer.render(data.path("userPrompt").asText(""), workspace.getVariables());

        if (userPrompt.isEmpty()) {
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.failure("LLM 节点的用户提示词为空").toMap());
        }

        ChatClient client = resolveChatClient(data);
        JsonNode outputsDecl = data.path("outputs");
        boolean structured = StructuredOutputHelper.needsStructuredOutput(data);

        log.info("[LLM] 节点 {} 开始调用, userPrompt长度={}, structured={}", nodeId, userPrompt.length(), structured);

        // 结构化模式：走 extractWithRetry 统一流程
        if (structured) {
            RetryPolicy policy = RetryPolicy.fromNodeData(data, workflowProperties);
            StructuredOutputHelper.ExtractionResult ext =
                    StructuredOutputHelper.extractWithRetry(client, userPrompt, systemPrompt, outputsDecl, policy);

            if (ext.errorMessage != null) {
                log.warn("[LLM] 节点 {} 结构化输出失败（重试 {} 次）: {}", nodeId, ext.attempts, ext.errorMessage);
                return NodeExecutionResult.success(nodeId,
                        NodeOutputEnvelope.failure(ext.errorMessage).toMap());
            }

            String msg = "LLM 调用成功" + (ext.attempts > 1 ? "（重试 " + (ext.attempts - 1) + " 次）" : "");
            log.info("[LLM] 节点 {} 完成, 输出字段数={}, attempts={}", nodeId, ext.data.size(), ext.attempts);
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.success(msg, ext.data).toMap());
        }

        // 单输出降级模式：未声明 outputs，按默认变量名写一个字段
        try {
            var request = client.prompt().user(userPrompt);
            if (!systemPrompt.isEmpty()) {
                request = request.system(systemPrompt);
            }
            String result = request.call().content();

            if (result == null || result.isBlank()) {
                return NodeExecutionResult.success(nodeId,
                        NodeOutputEnvelope.failure("LLM 返回空结果").toMap());
            }

            String varName = NodeTypeUtils.resolveOutputVarName(data, "result");
            Map<String, Object> dataMap = new LinkedHashMap<>();
            dataMap.put(varName, result);

            log.info("[LLM] 节点 {} 完成, 结果长度={}", nodeId, result.length());
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.success("LLM 调用成功", dataMap).toMap());

        } catch (Exception e) {
            log.error("[LLM] 节点 {} 调用失败: {}", nodeId, e.getMessage());
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.failure("LLM 调用失败: " + e.getMessage()).toMap());
        }
    }

    private ChatClient resolveChatClient(JsonNode data) {
        long modelId = data.path("modelId").asLong(0);
        if (modelId > 0) {
            try {
                return aiModelService.getChatClient(modelId);
            } catch (Exception e) {
                log.warn("[LLM] 指定模型 {} 获取失败，回退到默认模型: {}", modelId, e.getMessage());
            }
        }
        // 每次动态获取默认模型，确保模型配置变更能及时生效
        return aiModelService.getChatClientOrFallback(fallbackBuilder);
    }
}
