@Configuration
public class VectorStoreConfig {
    
    /**
     * 配置 Redis 连接池
     */
    @Bean
    public JedisPooled jedisPooled() {
        return new JedisPooled("127.0.0.1", 9379, null, "123456");
    }
    
    /**
     * 配置向量存储
     * 使用 Redis 作为向量数据库，Ollama 作为嵌入模型
     */
    @Bean
    public VectorStore vectorStore(
            @Qualifier("ollamaEmbeddingModel") EmbeddingModel embeddingModel,
            JedisPooled jedisPooled) {
        
        return RedisVectorStore.builder(jedisPooled, embeddingModel)
                .indexName("shotai-rag-index")      // 索引名称
                .prefix("shotai:doc:")              // Key 前缀
                .initializeSchema(true)             // 自动初始化 Schema
                .build();
    }
}

@Slf4j
@Service
public class DocumentServiceImpl implements DocumentService {
    
    @Autowired
    private VectorStore vectorStore;
    
    /**
     * 处理上传的文档
     * 1. 解析文档
     * 2. 分割文本
     * 3. 向量化
     * 4. 存储到 Redis
     */
    @Override
    public int processDocument(MultipartFile file) throws Exception {
        log.info("【文档处理】开始: {}, 大小: {} KB", 
                file.getOriginalFilename(), file.getSize() / 1024);
        
        // 1. 将 MultipartFile 转换为 Resource
        Resource resource = new ByteArrayResource(file.getBytes()) {
            @Override
            public String getFilename() {
                return file.getOriginalFilename();
            }
        };
        
        // 2. 读取文档
        List<Document> documents = readDocument(resource, file.getOriginalFilename());
        
        // 3. 分割文档为更小的片段
        TokenTextSplitter splitter = new TokenTextSplitter(
                500,   // 每个片段最大 token 数
                100,   // 片段之间的重叠 token 数
                5,     // 最小片段长度
                10000, // 最大片段长度
                true   // 保持分隔符
        );
        
        List<Document> splitDocuments = splitter.apply(documents);
        
        // 4. 添加元数据
        String fileName = file.getOriginalFilename();
        for (Document doc : splitDocuments) {
            doc.getMetadata().put("fileName", fileName);
            doc.getMetadata().put("uploadTime", System.currentTimeMillis());
        }
        
        // 5. 向量化并存储到 Redis
        vectorStore.add(splitDocuments);
        log.info("【文档处理】完成: {}, 生成 {} 个向量", 
                fileName, splitDocuments.size());
        
        return splitDocuments.size();
    }
}

/**
 * 根据文件类型选择合适的文档读取器
 */
private List<Document> readDocument(Resource resource, String filename) {
    String lowerCaseFilename = filename.toLowerCase();
    
    // PDF 文件使用专门的 PDF 读取器
    if (lowerCaseFilename.endsWith(".pdf")) {
        PdfDocumentReaderConfig config = PdfDocumentReaderConfig.builder()
                .withPageTopMargin(0)
                .withPageBottomMargin(0)
                .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                        .withNumberOfTopTextLinesToDelete(0)
                        .build())
                .withPagesPerDocument(1)  // 每页作为一个文档
                .build();
        
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(resource, config);
        return pdfReader.get();
    } 
    // 其他格式使用 Tika 读取器（支持 txt, doc, docx, html 等）
    else {
        TikaDocumentReader tikaReader = new TikaDocumentReader(resource);
        return tikaReader.get();
    }
}

@Slf4j
@Service
public class RagServiceImpl implements RagService {
    
    @Autowired
    private VectorStore vectorStore;
    
    @Autowired
    @Qualifier("deepSeekChatModel")
    private ChatModel chatModel;
    
    private ChatClient chatClient;
    
    /**
     * RAG 提示词模板
     */
    private static final String RAG_PROMPT_TEMPLATE = """
            你是一个智能助手。请基于以下提供的相关文档内容来回答用户的问题。
            
            相关文档内容：
            {context}
            
            用户问题：
            {question}
            
            请注意：
            1. 如果文档内容中包含答案，请基于文档内容进行详细回答
            2. 如果文档内容不足以回答问题，请说明文档中没有相关信息，但可以给出一些建议
            3. 保持回答的准确性和相关性
            4. 用中文回答
            
            回答：
            """;
    
    /**
     * 使用 RAG 进行流式对话
     */
    @Override
    public Flux<String> ragChatStream(String userId, String message) {
        // 1. 从向量数据库检索相关文档
        String context = retrieveRelevantDocuments(message, 5);
        
        if (context.isEmpty()) {
            log.warn("【RAG检索】未找到相关文档，用户: {}", userId);
            context = "暂无相关文档内容。";
        } else {
            log.info("【RAG检索】找到相关文档 {} 字符，用户: {}", 
                    context.length(), userId);
        }
        
        // 2. 构建 RAG 提示词
        PromptTemplate promptTemplate = new PromptTemplate(RAG_PROMPT_TEMPLATE);
        Prompt prompt = promptTemplate.create(Map.of(
                "context", context,
                "question", message
        ));
        
        // 3. 使用 ChatClient 生成流式响应
        return chatClient.prompt(prompt).stream().content();
    }
}

/**
 * 检索相关文档片段
 */
@Override
public String retrieveRelevantDocuments(String query, int topK) {
    try {
        // 从向量数据库检索最相关的文档
        SearchRequest searchRequest = SearchRequest.builder()
                .query(query)                    // 查询文本
                .topK(topK)                      // 返回前 K 个结果
                .similarityThreshold(0.5)        // 相似度阈值（0-1）
                .build();
        
        List<Document> relevantDocuments = vectorStore.similaritySearch(searchRequest);
        
        // 将文档内容合并为上下文
        return relevantDocuments.stream()
                .map(doc -> {
                    String content = doc.getText();
                    String fileName = (String) doc.getMetadata().get("fileName");
                    return String.format("【来源: %s】\n%s", fileName, content);
                })
                .collect(Collectors.joining("\n\n---\n\n"));
        
    } catch (Exception e) {
        log.error("【向量检索失败】错误: {}", e.getMessage());
        return "";
    }
}

@Slf4j
@RestController
@RequestMapping("/api/rag")
public class RagController {
    
    @Autowired
    private DocumentService documentService;
    
    /**
     * 上传文档
     */
    @PostMapping(value = "/upload", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    public DocumentUploadResponse uploadDocument(@RequestParam("file") MultipartFile file) {
        try {
            // 1. 验证文件
            if (file.isEmpty()) {
                return DocumentUploadResponse.error("上传的文件为空");
            }
            
            // 2. 验证文件类型
            if (!isSupportedFileType(file.getOriginalFilename())) {
                return DocumentUploadResponse.error(
                    "不支持的文件类型，目前支持: PDF, TXT, DOC, DOCX");
            }
            
            // 3. 处理文档（解析、分割、向量化、存储）
            int vectorCount = documentService.processDocument(file);
            
            // 4. 生成文档 ID
            String documentId = UUID.randomUUID().toString();
            
            log.info("文档处理成功: {}, 生成 {} 个向量", 
                    file.getOriginalFilename(), vectorCount);
            
            return DocumentUploadResponse.success(
                    documentId, 
                    file.getOriginalFilename(), 
                    vectorCount);
            
        } catch (Exception e) {
            log.error("处理文档上传时发生错误", e);
            return DocumentUploadResponse.error("文档处理失败: " + e.getMessage());
        }
    }
    
    /**
     * 检查文件类型是否支持
     */
    private boolean isSupportedFileType(String filename) {
        if (filename == null) return false;
        
        String lowerCaseFilename = filename.toLowerCase();
        return lowerCaseFilename.endsWith(".pdf") 
                || lowerCaseFilename.endsWith(".txt")
                || lowerCaseFilename.endsWith(".doc")
                || lowerCaseFilename.endsWith(".docx");
    }
}

// 余弦相似度公式
similarity = (A · B) / (||A|| * ||B||)

// 范围：-1 到 1
// 1: 完全相同
// 0: 无关
// -1: 完全相反

// 结合关键词检索和向量检索
public List<Document> hybridSearch(String query) {
    // 1. 向量检索
    List<Document> vectorResults = vectorStore.similaritySearch(query);
    
    // 2. 关键词检索（BM25）
    List<Document> keywordResults = keywordSearch(query);
    
    // 3. 合并结果，去重
    return mergeAndRank(vectorResults, keywordResults);
}

// 使用更强的模型对检索结果重新排序
public List<Document> rerank(List<Document> documents, String query) {
    return documents.stream()
            .sorted((d1, d2) -> {
                double score1 = rerankModel.score(query, d1.getText());
                double score2 = rerankModel.score(query, d2.getText());
                return Double.compare(score2, score1);
            })
            .collect(Collectors.toList());
}

@Cacheable(value = "rag-results", key = "#query")
public String retrieveRelevantDocuments(String query, int topK) {
    // 对于相同的问题，直接返回缓存结果
}

// 降低阈值
.similarityThreshold(0.3)

// 调整分割参数
new TokenTextSplitter(300, 50, ...)

// 增加检索数量
.topK(10)

// 优化提示词
"""
请严格基于以下文档回答，不要编造信息：
{context}
"""

// 限制文档大小
if (file.getSize() > 10 * 1024 * 1024) {
    throw new RuntimeException("文件太大，请上传小于 10MB 的文件");
}

// 使用索引加速
RedisVectorStore.builder()
    .indexName("shotai-rag-index")
    .build();