package com.agent.management.engine.executor;

import com.agent.management.engine.NodeExecutionResult;
import com.agent.management.engine.NodeExecutor;
import com.agent.management.engine.NodeOutputEnvelope;
import com.agent.management.engine.NodeWorkspace;
import com.agent.management.engine.TemplateRenderer;
import com.agent.management.engine.WorkflowContext;
import com.agent.management.model.entity.DataSource;
import com.agent.management.model.entity.GraphSource;
import com.agent.management.rag.document.DocumentRagRetriever;
import com.agent.management.rag.graph.GraphRagRetriever;
import com.agent.management.rag.kb.KnowledgeBaseRagRetriever;
import com.agent.management.rag.model.KnowledgeBaseRagRequest;
import com.agent.management.rag.model.KnowledgeBaseRagResult;
import com.agent.management.rag.model.RagEvidence;
import com.agent.management.rag.model.RagQuery;
import com.agent.management.rag.model.RagRetrievalResult;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.rag.structured.StructuredDataRagRetriever;
import com.agent.management.service.DataSourceService;
import com.agent.management.service.GraphSourceService;
import com.fasterxml.jackson.databind.JsonNode;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 知识库检索节点执行器
 *
 * <p>支持四种检索来源：</p>
 * <ul>
 *   <li><b>document</b>：文档数据库（DocumentRagRetriever）。sourceIds 支持多个文档 ID，
 *       allDocuments=true 时 sourceIds 留空，由 retriever 检索全部文档。</li>
 *   <li><b>structured</b>：结构化数据库（StructuredDataRagRetriever）。retriever 仅取 sourceIds[0]，
 *       因此多数据源时本执行器循环调用每个 datasourceId 并合并 evidences。</li>
 *   <li><b>graph</b>：知识图谱库（GraphRagRetriever）。retriever 仅取 sourceIds[0]，
 *       多图谱时同样循环调用合并。</li>
 *   <li><b>hybrid</b>：混合检索（KnowledgeBaseRagRetriever）。按知识库绑定关系自动并行调用上述三类 retriever。</li>
 * </ul>
 *
 * <p>query 字段支持 {@code {{变量名}}} 模板渲染，可引用上游节点输出。</p>
 *
 * <p>输出写入工作流上下文：</p>
 * <ul>
 *   <li><b>evidences</b>：{@code List<RagEvidence>}</li>
 *   <li><b>evidenceCount</b>：int</li>
 *   <li><b>sourceType</b>：DOCUMENT / STRUCTURED_DATA / GRAPH / HYBRID</li>
 *   <li><b>diagnostics</b>：{@code Map<String,Object>}</li>
 * </ul>
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class KnowledgeRetrievalExecutor implements NodeExecutor {

    private final DocumentRagRetriever documentRetriever;
    private final StructuredDataRagRetriever structuredRetriever;
    private final GraphRagRetriever graphRetriever;
    private final KnowledgeBaseRagRetriever knowledgeBaseRetriever;
    private final DataSourceService dataSourceService;
    private final GraphSourceService graphSourceService;

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

    @Override
    public NodeExecutionResult execute(String nodeId, JsonNode data, WorkflowContext context, NodeWorkspace workspace) {
        String source = data.path("source").asText("document");
        String rawQuery = data.path("query").asText("");
        String query = TemplateRenderer.render(rawQuery, workspace.getVariables()).trim();

        if (query.isEmpty()) {
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.failure("知识库检索节点的检索语句为空").toMap());
        }

        int topK = data.path("topK").asInt(5);
        if (topK <= 0) topK = 5;

        try {
            Map<String, Object> output = switch (source) {
                case "document" -> retrieveDocument(data, query, topK);
                case "structured" -> retrieveStructured(data, query, topK);
                case "graph" -> retrieveGraph(data, query, topK);
                case "hybrid" -> retrieveHybrid(data, query, topK);
                default -> {
                    log.warn("[KnowledgeRetrieval] 节点 {} 未知 source={}, 降级为 document", nodeId, source);
                    yield retrieveDocument(data, query, topK);
                }
            };

            log.info("[KnowledgeRetrieval] 节点 {} 完成, source={}, evidences={}",
                    nodeId, source, output.get("evidenceCount"));
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.success("知识检索完成", output).toMap());
        } catch (IllegalArgumentException e) {
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.failure(e.getMessage()).toMap());
        } catch (Exception e) {
            log.error("[KnowledgeRetrieval] 节点 {} 检索失败: {}", nodeId, e.getMessage(), e);
            return NodeExecutionResult.success(nodeId,
                    NodeOutputEnvelope.failure("知识库检索失败：" + e.getMessage()).toMap());
        }
    }

    // =========================== 检索实现 ===========================

    private Map<String, Object> retrieveDocument(JsonNode data, String query, int topK) {
        boolean allDocuments = data.path("allDocuments").asBoolean(true);
        List<String> documentIds = stringList(data, "documentIds");

        if (!allDocuments && documentIds.isEmpty()) {
            throw new IllegalArgumentException("未选择文档范围，请勾选文档或改为「在所有文档中检索」");
        }

        RagQuery rq = buildQuery(query, topK, allDocuments ? null : documentIds);
        rq.setFilters(Map.of("mode", "hybrid"));

        RagRetrievalResult result = documentRetriever.retrieve(rq);
        return toOutputMap(result);
    }

    private Map<String, Object> retrieveStructured(JsonNode data, String query, int topK) {
        List<String> datasourceIds = resolveScopedSourceIds(
                data, "allDatasources", "datasourceIds",
                () -> dataSourceService.listAll().stream().map(DataSource::getId).map(String::valueOf).toList());
        if (datasourceIds.isEmpty()) {
            throw new IllegalArgumentException("未选择结构化数据源，且系统中无可用数据源");
        }

        List<RagEvidence> evidences = new ArrayList<>();
        Map<String, Object> diagnostics = new LinkedHashMap<>();
        for (String dsId : datasourceIds) {
            RagQuery rq = buildQuery(query, topK, List.of(dsId));
            Map<String, Object> filters = new LinkedHashMap<>();
            filters.put("allowTextToSql", true);
            filters.put("maxRows", topK);
            rq.setFilters(filters);
            try {
                RagRetrievalResult result = structuredRetriever.retrieve(rq);
                mergeRetrieval(result, "datasource:" + dsId, evidences, diagnostics);
            } catch (Exception e) {
                diagnostics.put("error:datasource:" + dsId, e.getMessage());
                log.warn("[KnowledgeRetrieval] 结构化数据源 {} 检索失败: {}", dsId, e.getMessage());
            }
        }
        return aggregate(RagSourceType.STRUCTURED_DATA, evidences, diagnostics);
    }

    private Map<String, Object> retrieveGraph(JsonNode data, String query, int topK) {
        List<String> graphIds = resolveScopedSourceIds(
                data, "allGraphSources", "graphSourceIds",
                () -> graphSourceService.listAll().stream().map(GraphSource::getId).map(String::valueOf).toList());
        if (graphIds.isEmpty()) {
            throw new IllegalArgumentException("未选择图谱数据源，且系统中无可用图谱");
        }

        List<RagEvidence> evidences = new ArrayList<>();
        Map<String, Object> diagnostics = new LinkedHashMap<>();
        for (String gId : graphIds) {
            RagQuery rq = buildQuery(query, topK, List.of(gId));
            rq.setFilters(Map.of("allowTextToCypher", true));
            try {
                RagRetrievalResult result = graphRetriever.retrieve(rq);
                mergeRetrieval(result, "graph:" + gId, evidences, diagnostics);
            } catch (Exception e) {
                diagnostics.put("error:graph:" + gId, e.getMessage());
                log.warn("[KnowledgeRetrieval] 图谱 {} 检索失败: {}", gId, e.getMessage());
            }
        }
        return aggregate(RagSourceType.GRAPH, evidences, diagnostics);
    }

    private Map<String, Object> retrieveHybrid(JsonNode data, String query, int topK) {
        long kbId = data.path("knowledgeBaseId").asLong(0);
        if (kbId <= 0) {
            throw new IllegalArgumentException("混合检索需要指定有效的 knowledgeBaseId");
        }

        KnowledgeBaseRagRequest req = new KnowledgeBaseRagRequest();
        req.setKnowledgeBaseId(kbId);
        req.setQuery(query);
        req.setTopK(topK);
        KnowledgeBaseRagResult result = knowledgeBaseRetriever.retrieve(req);

        Map<String, Object> output = new LinkedHashMap<>();
        output.put("evidences", result.getEvidences());
        output.put("evidenceCount", result.getEvidences().size());
        output.put("sourceType", "HYBRID");
        output.put("knowledgeBaseId", kbId);
        output.put("enabledSources", result.getEnabledSources());
        output.put("diagnostics", result.getDiagnostics() == null ? Map.of() : result.getDiagnostics());
        return output;
    }

    // =========================== 工具 ===========================

    private static RagQuery buildQuery(String query, int topK, List<String> sourceIds) {
        RagQuery rq = new RagQuery();
        rq.setQuery(query);
        rq.setTopK(topK);
        if (sourceIds != null && !sourceIds.isEmpty()) rq.setSourceIds(sourceIds);
        return rq;
    }

    /**
     * 解析 sourceIds：若 allFlag=true，则调用 allSupplier 取全部；否则取显式 IDs。
     * 用于结构化和图谱的「在所有数据源中检索」语义。
     */
    private static List<String> resolveScopedSourceIds(JsonNode data, String allFlag, String idsField,
                                                        java.util.function.Supplier<List<String>> allSupplier) {
        boolean all = data.path(allFlag).asBoolean(false);
        if (all) return allSupplier.get();
        return stringList(data, idsField);
    }

    private static List<String> stringList(JsonNode data, String field) {
        JsonNode arr = data.path(field);
        if (!arr.isArray() || arr.isEmpty()) return List.of();
        List<String> result = new ArrayList<>();
        for (JsonNode item : arr) {
            String s = item.asText("");
            if (!s.isBlank()) result.add(s);
        }
        return result;
    }

    private static void mergeRetrieval(RagRetrievalResult result, String diagnosticKey,
                                        List<RagEvidence> evidences, Map<String, Object> diagnostics) {
        if (result == null) return;
        if (result.getEvidences() != null) evidences.addAll(result.getEvidences());
        if (result.getDiagnostics() != null && !result.getDiagnostics().isEmpty()) {
            diagnostics.put(diagnosticKey, result.getDiagnostics());
        }
    }

    private static Map<String, Object> toOutputMap(RagRetrievalResult result) {
        Map<String, Object> out = new LinkedHashMap<>();
        out.put("evidences", result.getEvidences() == null ? List.of() : result.getEvidences());
        out.put("evidenceCount", result.getEvidences() == null ? 0 : result.getEvidences().size());
        out.put("sourceType", result.getSourceType() == null ? "" : result.getSourceType().name());
        out.put("diagnostics", result.getDiagnostics() == null ? Map.of() : result.getDiagnostics());
        return out;
    }

    private static Map<String, Object> aggregate(RagSourceType type, List<RagEvidence> evidences,
                                                   Map<String, Object> diagnostics) {
        Map<String, Object> out = new LinkedHashMap<>();
        out.put("evidences", evidences);
        out.put("evidenceCount", evidences.size());
        out.put("sourceType", type.name());
        out.put("diagnostics", diagnostics);
        return out;
    }
}
