@Component
public class EntityRelationExtractor {

    private final ChatLanguageModel llm;

    public EntityRelationExtractor(ChatLanguageModel llm) {
        this.llm = llm;
    }

    /**
     * 从文本块中提取实体和关系
     * 核心思路：用结构化 prompt 让 LLM 返回 JSON 格式的三元组
     */
    public ExtractionResult extract(String text) {
        String prompt = """
            请从以下文本中提取实体和关系，严格按 JSON 格式返回。

            实体类型：人物、组织、技术、项目、概念、产品
            关系类型：使用、开发、依赖、属于、包含、合作、竞争

            文本：
            %s

            返回格式：
            {
              "entities": [
                {"name": "实体名", "type": "类型", "description": "一句话描述"}
              ],
              "relations": [
                {"source": "源实体", "target": "目标实体", "type": "关系类型"}
              ]
            }
            """.formatted(text);

        String response = llm.generate(prompt);
        return parseResponse(response);
    }

    private ExtractionResult parseResponse(String json) {
        // 解析 JSON，提取实体和关系
        // 实际生产中建议用 Jackson 或 Gson 做健壮解析
        ObjectMapper mapper = new ObjectMapper();
        try {
            return mapper.readValue(json, ExtractionResult.class);
        } catch (Exception e) {
            // LLM 返回的 JSON 可能格式不标准，需要容错处理
            String cleaned = json.replaceAll("```json\\s*", "").replaceAll("```", "").trim();
            return mapper.readValue(cleaned, ExtractionResult.class);
        }
    }
}

@Component
public class KnowledgeGraph {

    // 用 ConcurrentHashMap 做内存存储，生产环境建议用 Neo4j
    private final Map<String, EntityNode> entities = new ConcurrentHashMap<>();
    private final List<RelationEdge> relations = new CopyOnWriteArrayList<>();

    /**
     * 添加实体，如果已存在则合并描述
     */
    public void addEntity(String name, String type, String description) {
        entities.merge(name, new EntityNode(name, type, description),
            (existing, newNode) -> {
                // 实体消歧：合并描述信息
                existing.setDescription(existing.getDescription() + "; " + description);
                return existing;
            });
    }

    /**
     * 添加关系，自动创建缺失的实体节点
     */
    public void addRelation(String source, String target, String type) {
        // 确保两端实体存在
        entities.putIfAbsent(source, new EntityNode(source, "未知", ""));
        entities.putIfAbsent(target, new EntityNode(target, "未知", ""));

        // 避免重复关系
        boolean exists = relations.stream()
            .anyMatch(r -> r.getSource().equals(source)
                && r.getTarget().equals(target)
                && r.getType().equals(type));
        if (!exists) {
            relations.add(new RelationEdge(source, target, type));
        }
    }

    /**
     * 局部搜索：从指定实体出发，遍历 N 跳内的关联信息
     */
    public SubGraph localSearch(String entityName, int hops) {
        Set<String> visited = new HashSet<>();
        Set<String> currentLevel = new HashSet<>();
        currentLevel.add(entityName);
        List<EntityNode> subEntities = new ArrayList<>();
        List<RelationEdge> subRelations = new ArrayList<>();

        for (int hop = 0; hop < hops; hop++) {
            Set<String> nextLevel = new HashSet<>();
            for (String name : currentLevel) {
                if (visited.contains(name)) continue;
                visited.add(name);

                EntityNode entity = entities.get(name);
                if (entity != null) subEntities.add(entity);

                // 找到所有与该实体相关的关系
                for (RelationEdge rel : relations) {
                    if (rel.getSource().equals(name) || rel.getTarget().equals(name)) {
                        subRelations.add(rel);
                        nextLevel.add(rel.getSource());
                        nextLevel.add(rel.getTarget());
                    }
                }
            }
            currentLevel = nextLevel;
        }

        return new SubGraph(subEntities, subRelations);
    }
}

@Component
public class GlobalSearchEngine {

    private final ChatLanguageModel llm;
    private final KnowledgeGraph knowledgeGraph;

    public GlobalSearchEngine(ChatLanguageModel llm, KnowledgeGraph knowledgeGraph) {
        this.llm = llm;
        this.knowledgeGraph = knowledgeGraph;
    }

    /**
     * 全局搜索：对图谱中所有实体做 Map-Reduce 式摘要
     *
     * 完整版 GraphRAG 会先做社区检测，这里简化为按实体类型分组
     */
    public String globalSearch(String query) {
        // Map 阶段：按实体类型分组，每组生成中间答案
        Map<String, List<EntityNode>> byType = knowledgeGraph.getAllEntities()
            .stream()
            .collect(Collectors.groupingBy(EntityNode::getType));

        List<String> intermediateAnswers = new ArrayList<>();

        for (Map.Entry<String, List<EntityNode>> entry : byType.entrySet()) {
            String type = entry.getKey();
            List<EntityNode> entities = entry.getValue();

            // 拼接该类型所有实体的信息
            StringBuilder context = new StringBuilder();
            for (EntityNode entity : entities) {
                context.append("- ").append(entity.getName())
                       .append(": ").append(entity.getDescription()).append("\n");

                // 加入关系信息
                knowledgeGraph.getRelationsFor(entity.getName())
                    .forEach(rel -> context.append("  关系: ")
                        .append(rel.getSource()).append(" → ")
                        .append(rel.getType()).append(" → ")
                        .append(rel.getTarget()).append("\n"));
            }

            String mapPrompt = """
                基于以下%s类别的知识信息，回答问题：
                %s

                知识信息：
                %s

                请从%s的角度，给出简洁但全面的回答（200字以内）。
                """.formatted(type, query, context, type);

            String answer = llm.generate(mapPrompt);
            intermediateAnswers.add("【" + type + "视角】" + answer);
        }

        // Reduce 阶段：汇总所有中间答案
        String reducePrompt = """
            用户问题：%s

            以下是不同视角的分析结果，请综合归纳，给出一个完整、有层次的回答：

            %s

            要求：
            1. 提炼各视角的共性结论
            2. 指出不同视角之间的关联
            3. 给出全局性总结
            """.formatted(query, String.join("\n\n", intermediateAnswers));

        return llm.generate(reducePrompt);
    }
}

@Service
public class GraphRAGService {

    private final EntityRelationExtractor extractor;
    private final KnowledgeGraph knowledgeGraph;
    private final GlobalSearchEngine globalSearch;
    private final EmbeddingStore<TextSegment> embeddingStore;

    /**
     * 索引构建：文档 → 实体提取 → 图谱构建 → 向量化存储
     */
    public void buildIndex(List<Document> documents) {
        for (Document doc : documents) {
            // 1. 分块
            List<TextSegment> chunks = new DocumentSplitter()
                .split(doc, 500, 50); // 500 字一块，50 字重叠

            for (TextSegment chunk : chunks) {
                // 2. 实体-关系提取
                ExtractionResult result = extractor.extract(chunk.text());

                // 3. 写入知识图谱
                result.getEntities().forEach(e ->
                    knowledgeGraph.addEntity(e.getName(), e.getType(), e.getDescription()));
                result.getRelations().forEach(r ->
                    knowledgeGraph.addRelation(r.getSource(), r.getTarget(), r.getType()));

                // 4. 同时做向量化存储（用于局部搜索的文档检索）
                Embedding embedding = embeddingModel.embed(chunk).content();
                embeddingStore.add(embedding, chunk);
            }
        }
        log.info("索引构建完成。实体数：{}，关系数：{}",
            knowledgeGraph.getEntityCount(), knowledgeGraph.getRelationCount());
    }

    /**
     * 智能查询：根据问题类型自动选择局部或全局搜索
     */
    public String query(String question) {
        // 简单启发式：包含"整体"、"全局"、"总结"、"概述"等词时用全局搜索
        boolean isGlobal = containsGlobalKeywords(question);

        if (isGlobal) {
            return globalSearch.globalSearch(question);
        } else {
            // 局部搜索：先识别实体，再图谱遍历 + 向量检索
            String entity = extractMainEntity(question);
            SubGraph subGraph = knowledgeGraph.localSearch(entity, 2);
            List<TextSegment> docs = vectorSearch(question, 5);

            String context = buildContext(subGraph, docs);
            return llm.generate("基于以下信息回答问题：\n" + context + "\n\n问题：" + question);
        }
    }
}

/**
 * 批量实体提取：将多个文本块合并为一次 LLM 调用
 */
public List<ExtractionResult> batchExtract(List<TextSegment> chunks) {
    StringBuilder combined = new StringBuilder();
    for (int i = 0; i < chunks.size(); i++) {
        combined.append("--- 文档块 ").append(i + 1).append(" ---\n");
        combined.append(chunks.get(i).text()).append("\n\n");
    }

    String prompt = """
        请分别从以下 %d 个文档块中提取实体和关系。
        对每个文档块，返回独立的 JSON 对象。

        %s

        返回格式：JSON 数组，每个元素对应一个文档块的提取结果。
        """.formatted(chunks.size(), combined);

    return llm.generate(prompt); // 一次调用，批量提取
}

/**
 * 基于嵌入相似度的实体消歧
 */
public String resolveEntity(String newEntity, double threshold) {
    Embedding newEmb = embeddingModel.embed(newEntity).content();

    for (String existing : knowledgeGraph.getAllEntityNames()) {
        Embedding existingEmb = embeddingModel.embed(existing).content();
        double similarity = CosineSimilarity.between(newEmb, existingEmb);

        if (similarity > threshold) {
            return existing; // 返回已存在的规范名
        }
    }
    return newEntity; // 没找到匹配，作为新实体
}

/**
 * 将 GraphRAG 注册为 Agent 的工具
 * 参考：AI Agent 的工具箱（/posts/fd311830.html）
 */
@Tool("搜索知识库中的信息。对于全局性问题使用全局模式，精确查找使用局部模式。")
public String searchKnowledge(
    @P("查询内容") String query,
    @P("搜索模式：local 或 global") String mode
) {
    if ("global".equals(mode)) {
        return globalSearchEngine.globalSearch(query);
    } else {
        return localSearchEngine.localSearch(query);
    }
}