SuperBizAgent-java/src/main/java/com/superbiz/agent/service/KnowledgeBaseInitService.java

package com.superbiz.agent.service;

import com.superbiz.agent.domain.entity.ApiDocument;
import com.superbiz.agent.domain.enums.FaultCategory;
import com.superbiz.agent.repository.ApiDocumentRepository;
import com.superbiz.agent.dto.KnowledgeEntry;
import com.superbiz.agent.dto.Frontmatter;
import com.superbiz.agent.dto.DocumentChunk;
import lombok.Data;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.io.IOException;
import java.nio.file.*;
import java.nio.file.attribute.BasicFileAttributes;
import java.time.LocalDateTime;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Collectors;

/**
 * 知识库初始化服务
 * 负责批量导入 knowledge_base 目录下的文档到数据库和 Milvus
 */
@Service
public class KnowledgeBaseInitService {

    private static final Logger logger = LoggerFactory.getLogger(KnowledgeBaseInitService.class);

    @Value("${knowledge.base-path:knowledge_base}")
    private String knowledgeBasePath;

    @Autowired
    private ApiDocumentRepository apiDocumentRepository;

    @Autowired
    private FrontmatterParser frontmatterParser;

    @Autowired
    private DocumentChunkService documentChunkService;

    @Autowired
    private VectorIndexService vectorIndexService;

    @Autowired
    private VectorEmbeddingService vectorEmbeddingService;

    @Autowired
    private KnowledgeIndexService knowledgeIndexService;

    /**
     * 初始化知识库
     *
     * @param force 是否强制重新导入（跳过去重检查）
     * @return 初始化结果
     */
    @Transactional(rollbackFor = Exception.class)
    public InitResult initializeKnowledgeBase(boolean force) {
        logger.info("开始初始化知识库: basePath={}, force={}", knowledgeBasePath, force);

        InitResult result = new InitResult();
        Path baseDir = Paths.get(knowledgeBasePath);

        if (!Files.exists(baseDir)) {
            logger.error("知识库目录不存在: {}", knowledgeBasePath);
            throw new RuntimeException("知识库目录不存在: " + knowledgeBasePath);
        }

        // 1. 扫描所有 Markdown 文件
        List<Path> markdownFiles = scanMarkdownFiles(baseDir);
        result.setScanned(markdownFiles.size());
        logger.info("扫描到 {} 个 Markdown 文件", markdownFiles.size());

        // 2. 如果非强制模式，获取已存在的文档（用于去重）
        Set<String> existingFilePaths = new HashSet<>();
        if (!force) {
            existingFilePaths = apiDocumentRepository.findAll().stream()
                    .map(ApiDocument::getFilePath)
                    .collect(Collectors.toSet());
            logger.info("已存在  个文档记录", existingFilePaths.size());
        }

        // 3. 逐个处理文档
        for (Path file : markdownFiles) {
            String relativePath = baseDir.relativize(file).toString().replace("\\", "/");

            try {
                // 去重检查
                if (!force && existingFilePaths.contains(relativePath)) {
                    logger.debug("跳过已存在的文档: {}", relativePath);
                    result.incrementSkipped();
                    result.addDetail(relativePath, "已存在，跳过");
                    continue;
                }

                // 解析文档
                String content = Files.readString(file);
                Frontmatter frontmatter = frontmatterParser.parse(content);

                if (frontmatter == null) {
                    logger.warn("文档格式无效: {}, frontmatter 解析失败", relativePath);
                    result.incrementFailed();
                    result.addDetail(relativePath, "格式无效: frontmatter 解析失败");
                    continue;
                }

                // 提取字段
                String title = frontmatter.getTitle();
                String summary = frontmatter.getSummary();
                String category = frontmatter.getCategory() != null ? frontmatter.getCategory() : "general";
                List<String> keywords = frontmatter.getKeywords();

                if (title == null || title.isBlank()) {
                    logger.warn("文档缺少标题: {}", relativePath);
                    result.incrementFailed();
                    result.addDetail(relativePath, "缺少标题");
                    continue;
                }

                // 保存到数据库
                ApiDocument document = saveToDatabase(relativePath, title, summary, category, content, keywords);

                // 提取文档正文（去除 frontmatter）
                String body = extractBody(content);

                // 文档分块
                List<DocumentChunk> chunks = documentChunkService.chunkDocument(body, relativePath);
                logger.debug("文档分块完成: {} -> {} 个 chunk", relativePath, chunks.size());

                // 上传到 Milvus
                try {
                    vectorIndexService.indexDocumentChunks(document.getDocId(), chunks, category);

                    document.setStatus("INDEXED");
                    document.setChunkCount(chunks.size());
                    document.setIndexedAt(LocalDateTime.now());
                    apiDocumentRepository.save(document);

                    logger.info("文档已索引到 Milvus: {} (docId={}, chunks={})",
                            title, document.getDocId(), chunks.size());
                } catch (Exception e) {
                    logger.error("上传到 Milvus 失败: {}", relativePath, e);

                    document.setStatus("FAILED");
                    document.setErrorMessage(e.getMessage());
                    apiDocumentRepository.save(document);

                    result.incrementFailed();
                    result.addDetail(relativePath, "Milvus 索引失败: " + e.getMessage());
                    continue;  // 跳过该文档，继续处理下一个
                }

                // 添加到 L0 内存索引
                KnowledgeEntry entry = KnowledgeEntry.builder()
                        .filePath(relativePath)
                        .title(title)
                        .keywords(keywords)
                        .summary(summary)
                        .category(category)
                        .build();
                knowledgeIndexService.addToIndex(entry);

                result.incrementInserted();
                result.addDetail(relativePath, "导入成功（L0+L1）");
                logger.info("文档导入成功: {} -> {} (L0+L1 索引已更新)", relativePath, title);

            } catch (Exception e) {
                logger.error("处理文档失败: {}", relativePath, e);
                result.incrementFailed();
                result.addDetail(relativePath, "处理失败: " + e.getMessage());
            }
        }

        logger.info("知识库初始化完成: 扫描={}, 跳过={}, 新增={}, 失败={}",
                result.getScanned(), result.getSkipped(), result.getInserted(), result.getFailed());

        return result;
    }

    /**
     * 获取知识库统计信息
     */
    public Stats getStats() {
        Stats stats = new Stats();

        // 数据库中的文档数量
        long totalDocuments = apiDocumentRepository.count();
        stats.setTotalDocuments(totalDocuments);

        // L0 索引中的文档数量
        int indexSize = knowledgeIndexService.getIndexSize();
        logger.debug("L0 索引大小: {}", indexSize);

        // 按分类统计（从 fault_category 字段读取）
        Map<String, Long> categoryCount = apiDocumentRepository.findAll().stream()
                .collect(Collectors.groupingBy(
                        doc -> doc.getFaultCategory() != null ? doc.getFaultCategory().name() : "GENERAL",
                        Collectors.counting()
                ));
        stats.setCategoryCount(categoryCount);

        // Milvus 中的向量数量（需要实现）
        // TODO: 查询 Milvus collection 的实体数量
        stats.setTotalVectors(0L);

        return stats;
    }

    /**
     * 扫描目录下所有 Markdown 文件
     */
    private List<Path> scanMarkdownFiles(Path baseDir) {
        List<Path> files = new ArrayList<>();

        try {
            Files.walkFileTree(baseDir, new SimpleFileVisitor<Path>() {
                @Override
                public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) {
                    if (file.toString().endsWith(".md")) {
                        files.add(file);
                    }
                    return FileVisitResult.CONTINUE;
                }

                @Override
                public FileVisitResult visitFileFailed(Path file, IOException exc) {
                    logger.warn("访问文件失败: {}", file, exc);
                    return FileVisitResult.CONTINUE;
                }
            });
        } catch (IOException e) {
            logger.error("扫描目录失败: {}", baseDir, e);
            throw new RuntimeException("扫描目录失败", e);
        }

        return files;
    }

    /**
     * 保存文档到数据库
     */
    private ApiDocument saveToDatabase(String filePath, String title, String summary,
                                       String category, String content, List<String> keywords) {
        ApiDocument document = new ApiDocument();
        document.setDocId(UUID.randomUUID().toString());
        document.setFileName(Paths.get(filePath).getFileName().toString());
        document.setFilePath(filePath);
        document.setApiName(title);  // 使用 title 作为 apiName
        document.setStatus("PENDING");  // 初始状态为 PENDING，索引成功后更新为 INDEXED

        // 映射 category 到 FaultCategory 枚举
        FaultCategory faultCategory = FaultCategory.fromString(category);
        document.setFaultCategory(faultCategory);

        // 将 frontmatter 信息保存到 metadata（JSON 格式）
        String metadataJson = String.format(
                "{\"title\":\"%s\",\"summary\":\"%s\",\"category\":\"%s\",\"keywords\":%s}",
                escapeJson(title),
                escapeJson(summary),
                escapeJson(category),
                "[\"" + String.join("\",\"", keywords.stream().map(this::escapeJson).toArray(String[]::new)) + "\"]"
        );
        document.setMetadata(metadataJson);

        document.setFileSize((long) content.length());

        return apiDocumentRepository.save(document);
    }

    /**
     * JSON 转义
     */
    private String escapeJson(String str) {
        if (str == null) {
            return "";
        }
        return str.replace("\\", "\\\\")
                  .replace("\"", "\\\"")
                  .replace("\n", "\\n")
                  .replace("\r", "\\r");
    }

    /**
     * 提取文档正文（去除 frontmatter）
     */
    private String extractBody(String content) {
        if (!content.trim().startsWith("---")) {
            return content;
        }

        int firstEnd = content.indexOf("---", 3);
        if (firstEnd == -1) {
            return content;
        }

        int secondEnd = content.indexOf("---", firstEnd + 3);
        if (secondEnd == -1) {
            return content.substring(firstEnd + 3).trim();
        }

        return content.substring(secondEnd + 3).trim();
    }

    // ==================== 数据模型 ====================

    /**
     * 初始化结果
     */
    @Data
    public static class InitResult {
        private int scanned;          // 扫描到的文件数量
        private int skipped;          // 跳过的文件数量（已存在）
        private int inserted;         // 成功导入的文件数量
        private int failed;           // 失败的文件数量
        private Map<String, String> details = new LinkedHashMap<>();  // 详细信息

        public void incrementSkipped() {
            this.skipped++;
        }

        public void incrementInserted() {
            this.inserted++;
        }

        public void incrementFailed() {
            this.failed++;
        }

        public void addDetail(String filePath, String message) {
            this.details.put(filePath, message);
        }
    }

    /**
     * 统计信息
     */
    @Data
    public static class Stats {
        private long totalDocuments;              // 数据库中的文档总数
        private long totalVectors;                // Milvus 中的向量总数
        private Map<String, Long> categoryCount;  // 按分类统计
    }
}