package com.ruoyi.llm.service.impl;

import com.alibaba.fastjson2.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.ruoyi.common.config.RuoYiConfig;
import com.ruoyi.llm.domain.CmcChat;
import com.ruoyi.llm.domain.CmcDocument;
import com.ruoyi.common.utils.milvus.MilvusConnectionPool;
import com.ruoyi.llm.domain.AnswerStreamResponse;
import com.ruoyi.llm.service.ICmcChatService;
import com.ruoyi.llm.service.ICmcDocumentService;
import com.ruoyi.llm.service.ILangChainMilvusService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallzhv15.BgeSmallZhV15EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.collection.request.ReleaseCollectionReq;
import io.milvus.v2.service.vector.request.DeleteReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.SearchReq.SearchReqBuilder;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.usermodel.Paragraph;
import org.apache.poi.hwpf.usermodel.Range;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;
import org.apache.poi.xwpf.usermodel.XWPFStyle;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatSession;
import org.noear.solon.ai.chat.message.ChatMessage;
import org.noear.solon.ai.chat.prompt.Prompt;
import org.noear.solon.ai.chat.session.InMemoryChatSession;
import org.openxmlformats.schemas.wordprocessingml.x2006.main.CTPPrGeneral;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import javax.annotation.PostConstruct;
import java.io.*;
import java.util.*;

@Service
public class LangChainMilvusServiceImpl implements ILangChainMilvusService {
    @Autowired
    private ICmcChatService cmcChatService;

    @Autowired
    private ICmcDocumentService cmcDocumentService;

    private static final EmbeddingModel embeddingModel = new BgeSmallZhV15EmbeddingModel();

    private String processValue = "";

    @Value("${cmc.llmService.url}")
    private String llmServiceUrl;

    @Autowired
    private MilvusConnectionPool milvusConnectionPool;

    // 复用 ChatModel 实例
    private ChatModel chatModel;

    @PostConstruct
    public void initChatModel() {
        // 初始化 ChatModel
        chatModel = ChatModel.of(llmServiceUrl)
                .model("Qwen")
                .build();
    }

    /**
     * 上传多文件
     *
     * @return 结果
     */
    public String getProcess() {
        return processValue;
    }

    /**
     * 导入知识库文件
     * 
     * @return
     */
    @Override
    public int insertLangchainEmbeddingDocument(MultipartFile[] fileList, String collectionName) {
        processValue = "";
        int successfullyInsertedFiles = 0;

        for (int i = 0; i < fileList.length; i++) {
            MultipartFile file = fileList[i];
            try {
                // 构建上传目录路径
                String uploadDir = RuoYiConfig.getProfile() + "/upload/rag/knowledge/" + collectionName;
                File profilePath = new File(uploadDir);

                // 确保目录存在，创建目录并设置权限
                if (!profilePath.exists())
                    profilePath.mkdirs();
                File transferFile = new File(profilePath, file.getOriginalFilename());
                if (!transferFile.exists())
                    file.transferTo(transferFile);
                List<TextSegment> segments = splitDocument(transferFile);

                // 准备导入数据
                List<JsonObject> data = new ArrayList<>();
                // 添加一条包含所有二、三级标题的记录
                String titlesContent = extractAllTitles(transferFile);
                String[] fileName = file.getOriginalFilename().split("\\.");
                if (!titlesContent.trim().isEmpty()) {
                    JSONObject titlesJsonObj = new JSONObject();
                    titlesJsonObj.put("file_name", file.getOriginalFilename());
                    titlesJsonObj.put("file_type", fileName[fileName.length - 1]);
                    titlesJsonObj.put("title", "目录");
                    titlesJsonObj.put("content", titlesContent);
                    titlesJsonObj.put("embedding", embeddingModel.embed(titlesContent).content().vectorAsList());
                    String titlesJsonString = titlesJsonObj.toJSONString();
                    JsonObject titlesJsonObject = JsonParser.parseString(titlesJsonString).getAsJsonObject();
                    data.add(titlesJsonObject);
                }

                // 提取文本和生成嵌入
                for (TextSegment segment : segments) {
                    String text = segment.text();
                    if (text.trim().isEmpty())
                        continue;

                    JSONObject fastjsonObj = new JSONObject();
                    fastjsonObj.put("file_name", file.getOriginalFilename());
                    fastjsonObj.put("file_type", fileName[fileName.length - 1]);
                    fastjsonObj.put("title", text.split("\n")[0].split(" ")[1]);
                    fastjsonObj.put("content", text);
                    fastjsonObj.put("embedding", embeddingModel.embed(text).content().vectorAsList());
                    String jsonString = fastjsonObj.toJSONString();
                    JsonObject jsonObject = JsonParser.parseString(jsonString).getAsJsonObject();
                    data.add(jsonObject);
                }

                // 从连接池获取客户端
                MilvusClientV2 milvusClient = null;
                try {
                    milvusClient = milvusConnectionPool.getClient();
                    // 加载集合
                    LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                            .collectionName(collectionName)
                            .build();
                    milvusClient.loadCollection(loadCollectionReq);

                    try {
                        // 先删除相同文件名的记录
                        DeleteReq deleteReq = DeleteReq.builder()
                                .collectionName(collectionName)
                                .filter(String.format("file_name == \"%s\"", file.getOriginalFilename()))
                                .build();
                        if (milvusClient.delete(deleteReq).getDeleteCnt() > 0)
                            System.out.println("已删除同名文件记录: " + file.getOriginalFilename());

                        // 构建导入请求
                        InsertReq insertReq = InsertReq.builder()
                                .collectionName(collectionName)
                                .data(data)
                                .build();

                        // 执行导入并检查结果
                        InsertResp insertResp = milvusClient.insert(insertReq);
                        if (insertResp != null && insertResp.getInsertCnt() > 0) {
                            successfullyInsertedFiles++;
                            processValue = "上传中: " + String.format("%.2f", (double) (i + 1) / fileList.length * 100) + "%";
                            System.out.println("成功导入" + (i + 1) + "/" + fileList.length + "文件: "
                                    + file.getOriginalFilename() + ", 导入记录数: " + insertResp.getInsertCnt());
                        } else {
                            System.err.println("文件导入失败: " + file.getOriginalFilename() + ", 没有记录被导入");
                        }
                    } finally {
                        // 释放集合
                        ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                                .collectionName(collectionName)
                                .build();
                        milvusClient.releaseCollection(releaseCollectionReq);
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw new RuntimeException("获取Milvus连接失败", e);
                } finally {
                    // 将客户端返回到连接池
                    if (milvusClient != null) {
                        milvusConnectionPool.returnClient(milvusClient);
                    }
                }
            } catch (Exception e) {
                System.err.println("处理文件时出错: " + file.getOriginalFilename() + ", 错误: " + e.getMessage());
                e.printStackTrace();
                // 继续处理下一个文件
            }
        }

        processValue = "上传完成";
        return successfullyInsertedFiles;
    }

    /**
     * 从Milvus检索相关文档
     * 
     * @return
     */
    @Override
    public List<JSONObject> retrieveFromMilvus(String collectionName, String query, int topK) {
        List<JSONObject> resultList = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResultList = retrieve(collectionName, query, topK);
        searchResultList.forEach(searchResult -> {
            JSONObject result = new JSONObject();
            result.put("file_name", searchResult.get(0).getEntity().get("file_name"));
            result.put("content", searchResult.get(0).getEntity().get("content"));
            resultList.add(result);
        });
        return resultList;
    }

    /**
     * 从Milvus检索相关文档及相关度
     * 
     * @return
     */
    @Override
    public List<JSONObject> similarityFromMilvus(String collectionName, String query, int topK) {
        List<JSONObject> resultList = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResultList = retrieve(collectionName, query, topK);
        searchResultList.forEach(searchResult -> {
            JSONObject result = new JSONObject();
            result.put("score", searchResult.get(0).getScore());
            result.put("file_name", searchResult.get(0).getEntity().get("file_name"));
            result.put("content", searchResult.get(0).getEntity().get("content"));
            resultList.add(result);
        });
        resultList.removeIf(jsonObject -> jsonObject.getDouble("score") < 0.7);
        return resultList;
    }

    /**
     * 调用LLM生成回答
     * 
     * @return
     */
    @Override
    public SseEmitter generateAnswer(String topicId, String prompt) {
        // 设置SSE超时时间为10分钟
        SseEmitter emitter = new SseEmitter(600000L);

        // 设置回调
        emitter.onCompletion(() -> {
        });
        emitter.onTimeout(() -> {
            emitter.complete();
        });
        emitter.onError(e -> System.err.println("SSE连接错误: " + e.getMessage()));

        // 使用独立线程池处理，避免阻塞主线程
        ExecutorService executor = Executors.newSingleThreadExecutor();

        executor.execute(() -> {
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                // 用于累计token用量
                AtomicLong[] tokenUsage = new AtomicLong[3];
                tokenUsage[0] = new AtomicLong(0); // promptTokens
                tokenUsage[1] = new AtomicLong(0); // completionTokens
                tokenUsage[2] = new AtomicLong(0); // totalTokens

                List<ChatMessage> messages = new ArrayList<>();
                if (topicId != null) {
                    CmcChat cmcChat = new CmcChat();
                    cmcChat.setTopicId(topicId);
                    List<CmcChat> cmcChatList = cmcChatService.selectCmcChatList(cmcChat);
                    for (CmcChat chat : cmcChatList) {
                        messages.add(ChatMessage.ofUser(chat.getInput()));
                        messages.add(ChatMessage.ofAssistant(chat.getOutput()));
                    }
                }
                messages.add(ChatMessage.ofUser(prompt));
                ChatSession chatSession = InMemoryChatSession.builder().messages(messages).build();
                Prompt prompt1 = Prompt.of(prompt).attrPut("session", chatSession);

                // 估算prompt的token数量（包含历史对话）
                StringBuilder fullPrompt = new StringBuilder();
                for (ChatMessage msg : messages) {
                    fullPrompt.append(msg.getContent());
                }
                tokenUsage[0].set((long) fullPrompt.length());

                // 使用 CountDownLatch 等待生成完成
                CountDownLatch latch = new CountDownLatch(1);
                AtomicBoolean cancelled = new AtomicBoolean(false);

                chatModel.prompt(prompt1).stream().map(response -> {
                    String content = response.lastChoice().getMessage().getContent();
                    // 累计completion tokens
                    if (content != null) {
                        tokenUsage[1].addAndGet((long) content.length());
                        tokenUsage[2].set(tokenUsage[0].get() + tokenUsage[1].get());
                    }
                    return content;
                }).subscribe(
                        content -> {
                            if (cancelled.get()) {
                                return;
                            }
                            try {
                                if (content != null && !content.isEmpty()) {
                                    AnswerStreamResponse response = AnswerStreamResponse.content(content, false);
                                    response.setPromptTokens(tokenUsage[0].get());
                                    response.setCompletionTokens(tokenUsage[1].get());
                                    response.setTotalTokens(tokenUsage[2].get());
                                    emitter.send(SseEmitter.event()
                                            .name("answer")
                                            .data(objectMapper.writeValueAsString(response)));
                                }
                            } catch (Exception e) {
                                cancelled.set(true);
                            }
                        },
                        error -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                AnswerStreamResponse errorResponse = AnswerStreamResponse
                                        .error("生成回答时发生错误: " + error.getMessage());
                                emitter.send(SseEmitter.event()
                                        .name("error")
                                        .data(objectMapper.writeValueAsString(errorResponse)));
                            } catch (Exception sendEx) {
                            }
                            latch.countDown();
                        },
                        () -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                // 发送完成标记（包含最终token用量）
                                AnswerStreamResponse completeResponse = AnswerStreamResponse.content("[DONE]", true);
                                completeResponse.setPromptTokens(tokenUsage[0].get());
                                completeResponse.setCompletionTokens(tokenUsage[1].get());
                                completeResponse.setTotalTokens(tokenUsage[2].get());
                                emitter.send(SseEmitter.event()
                                        .name("completed")
                                        .data(objectMapper.writeValueAsString(completeResponse)));
                            } catch (Exception e) {
                            } finally {
                                latch.countDown();
                            }
                        });

                // 等待生成完成
                latch.await();

                emitter.complete();

            } catch (Exception e) {
                try {
                    AnswerStreamResponse errorResponse = AnswerStreamResponse.error("生成过程中发生错误: " + e.getMessage());
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data(objectMapper.writeValueAsString(errorResponse)));
                } catch (IOException sendEx) {
                    System.err.println("发送错误消息失败: " + sendEx.getMessage());
                }
                emitter.completeWithError(e);
            } finally {
                executor.shutdown();
            }
        });

        return emitter;
    }

    /**
     * 调用LLM+RAG生成回答
     * 
     * @return
     */
    @Override
    public SseEmitter generateAnswerWithCollection(String topicId, String question, List<JSONObject> contexts) {
        // 设置SSE超时时间为10分钟
        SseEmitter emitter = new SseEmitter(600000L);

        // 设置回调
        emitter.onCompletion(() -> {
        });
        emitter.onTimeout(() -> {
            emitter.complete();
        });
        emitter.onError(e -> System.err.println("SSE连接错误: " + e.getMessage()));

        // 使用独立线程池处理，避免阻塞主线程
        ExecutorService executor = Executors.newSingleThreadExecutor();

        executor.execute(() -> {
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                StringBuilder sb = new StringBuilder();
                sb.append("问题: ").append(question).append("\n\n");
                sb.append("根据以下上下文回答问题：\n\n");
                for (JSONObject context : contexts) {
                    sb.append("文件").append(": ")
                            .append(context.getString("file_name")).append("\n\n")
                            .append("上下文").append(": ")
                            .append(context.getString("content")).append("\n\n");
                }
                String prompt = sb.toString();

                // 用于累计token用量
                AtomicLong[] tokenUsage = new AtomicLong[3];
                tokenUsage[0] = new AtomicLong(0); // promptTokens
                tokenUsage[1] = new AtomicLong(0); // completionTokens
                tokenUsage[2] = new AtomicLong(0); // totalTokens

                // 使用 CountDownLatch 等待生成完成
                CountDownLatch latch = new CountDownLatch(1);
                AtomicBoolean cancelled = new AtomicBoolean(false);

                generateAnswerFluxWithUsage(prompt, tokenUsage).subscribe(
                        content -> {
                            if (cancelled.get()) {
                                return;
                            }
                            try {
                                if (content != null && !content.isEmpty()) {
                                    // 构建响应对象，包含token用量
                                    AnswerStreamResponse response = AnswerStreamResponse.content(content, false);
                                    response.setPromptTokens(tokenUsage[0].get());
                                    response.setCompletionTokens(tokenUsage[1].get());
                                    response.setTotalTokens(tokenUsage[2].get());
                                    emitter.send(SseEmitter.event()
                                            .name("answer")
                                            .data(objectMapper.writeValueAsString(response)));
                                }
                            } catch (Exception e) {
                                cancelled.set(true);
                            }
                        },
                        error -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                AnswerStreamResponse errorResponse = AnswerStreamResponse
                                        .error("生成回答时发生错误: " + error.getMessage());
                                emitter.send(SseEmitter.event()
                                        .name("error")
                                        .data(objectMapper.writeValueAsString(errorResponse)));
                            } catch (Exception sendEx) {
                            }
                            latch.countDown();
                        },
                        () -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                // 发送完成标记（包含最终token用量）
                                AnswerStreamResponse completeResponse = AnswerStreamResponse.content("[DONE]", true);
                                completeResponse.setPromptTokens(tokenUsage[0].get());
                                completeResponse.setCompletionTokens(tokenUsage[1].get());
                                completeResponse.setTotalTokens(tokenUsage[2].get());
                                emitter.send(SseEmitter.event()
                                        .name("completed")
                                        .data(objectMapper.writeValueAsString(completeResponse)));
                            } catch (Exception e) {
                            } finally {
                                latch.countDown();
                            }
                        });

                // 等待生成完成
                latch.await();

                emitter.complete();

            } catch (Exception e) {
                try {
                    AnswerStreamResponse errorResponse = AnswerStreamResponse.error("生成过程中发生错误: " + e.getMessage());
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data(objectMapper.writeValueAsString(errorResponse)));
                } catch (IOException sendEx) {
                    System.err.println("发送错误消息失败: " + sendEx.getMessage());
                }
                emitter.completeWithError(e);
            } finally {
                executor.shutdown();
            }
        });

        return emitter;
    }

    /**
     * 调用LLM生成回答（带token用量统计）
     * 
     * @param prompt     提示词
     * @param tokenUsage 用于收集token用量的数组 [promptTokens, completionTokens,
     *                   totalTokens]
     * @return Flux<String> 流式输出（包含内容）
     */
    private Flux<String> generateAnswerFluxWithUsage(String prompt, AtomicLong[] tokenUsage) {
        // 估算prompt的token数量（中文大约1个字符=1个token）
        if (tokenUsage != null && tokenUsage.length >= 3 && tokenUsage[0].get() == 0) {
            tokenUsage[0].set((long) prompt.length());
        }

        List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.ofUser(prompt));
        ChatSession chatSession = InMemoryChatSession.builder().messages(messages).build();
        Prompt prompt1 = Prompt.of(prompt).attrPut("session", chatSession);

        return chatModel.prompt(prompt1).stream()
                .map(response -> {
                    String content = response.lastChoice().getMessage().getContent();
                    // 累计completion tokens（中文大约1个字符=1个token）
                    if (tokenUsage != null && tokenUsage.length >= 3 && content != null) {
                        tokenUsage[1].addAndGet((long) content.length());
                        tokenUsage[2].set(tokenUsage[0].get() + tokenUsage[1].get());
                    }
                    return content;
                });
    }

    /**
     * 调用LLM生成回答
     */
    @Override
    public SseEmitter generateAnswerWithDocument(String topicId, String chatId, String question) throws Exception {
        // 设置SSE超时时间为10分钟
        SseEmitter emitter = new SseEmitter(600000L);

        // 设置回调
        emitter.onCompletion(() -> {
        });
        emitter.onTimeout(() -> {
            emitter.complete();
        });
        emitter.onError(e -> System.err.println("SSE连接错误: " + e.getMessage()));

        // 使用独立线程池处理，避免阻塞主线程
        ExecutorService executor = Executors.newSingleThreadExecutor();

        executor.execute(() -> {
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                // 用于累计token用量
                AtomicLong[] tokenUsage = new AtomicLong[3];
                tokenUsage[0] = new AtomicLong(0); // promptTokens
                tokenUsage[1] = new AtomicLong(0); // completionTokens
                tokenUsage[2] = new AtomicLong(0); // totalTokens

                CmcDocument cmcDocument = new CmcDocument();
                cmcDocument.setChatId(chatId);
                List<CmcDocument> documentList = cmcDocumentService.selectCmcDocumentList(cmcDocument);
                StringBuilder sb = new StringBuilder("问题: " + question + "\n\n").append("根据以下上下文回答问题：\n\n");
                for (CmcDocument document : documentList) {
                    File profilePath = new File(RuoYiConfig.getProfile() + document.getPath());
                    InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
                    List<TextSegment> segments = splitDocument(profilePath);
                    List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
                    embeddingStore.addAll(embeddings, segments);
                    Embedding queryEmbedding = embeddingModel.embed(question).content();
                    EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                            .queryEmbedding(queryEmbedding)
                            .minScore(0.7)
                            .build();
                    List<EmbeddingMatch<TextSegment>> results = embeddingStore.search(embeddingSearchRequest).matches();
                    results.sort(Comparator.comparingDouble(EmbeddingMatch<TextSegment>::score).reversed());
                    for (EmbeddingMatch<TextSegment> embeddingMatch : results) {
                        String contexts = embeddingMatch.embedded().toString();
                        sb.append("文件").append(": ")
                                .append(document.getPath()).append("\n\n")
                                .append("上下文").append(": ")
                                .append(contexts).append("\n\n");
                    }
                }

                String prompt = sb.toString();

                List<ChatMessage> messages = new ArrayList<>();
                if (topicId != null) {
                    CmcChat cmcChat = new CmcChat();
                    cmcChat.setTopicId(topicId);
                    List<CmcChat> cmcChatList = cmcChatService.selectCmcChatList(cmcChat);
                    for (CmcChat chat : cmcChatList) {
                        messages.add(ChatMessage.ofUser(chat.getInput()));
                        messages.add(ChatMessage.ofAssistant(chat.getOutput()));
                    }
                }
                messages.add(ChatMessage.ofUser(prompt));
                ChatSession chatSession = InMemoryChatSession.builder().messages(messages).build();
                Prompt prompt1 = Prompt.of(prompt).attrPut("session", chatSession);

                // 估算prompt的token数量（包含历史对话）
                StringBuilder fullPrompt = new StringBuilder();
                for (ChatMessage msg : messages) {
                    fullPrompt.append(msg.getContent());
                }
                tokenUsage[0].set((long) fullPrompt.length());

                // 使用 CountDownLatch 等待生成完成
                CountDownLatch latch = new CountDownLatch(1);
                AtomicBoolean cancelled = new AtomicBoolean(false);

                chatModel.prompt(prompt1).stream().map(response -> {
                    String content = response.lastChoice().getMessage().getContent();
                    // 累计completion tokens
                    if (content != null) {
                        tokenUsage[1].addAndGet((long) content.length());
                        tokenUsage[2].set(tokenUsage[0].get() + tokenUsage[1].get());
                    }
                    return content;
                }).subscribe(
                        content -> {
                            if (cancelled.get()) {
                                return;
                            }
                            try {
                                if (content != null && !content.isEmpty()) {
                                    AnswerStreamResponse response = AnswerStreamResponse.content(content, false);
                                    response.setPromptTokens(tokenUsage[0].get());
                                    response.setCompletionTokens(tokenUsage[1].get());
                                    response.setTotalTokens(tokenUsage[2].get());
                                    emitter.send(SseEmitter.event()
                                            .name("answer")
                                            .data(objectMapper.writeValueAsString(response)));
                                }
                            } catch (Exception e) {
                                cancelled.set(true);
                            }
                        },
                        error -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                AnswerStreamResponse errorResponse = AnswerStreamResponse
                                        .error("生成回答时发生错误: " + error.getMessage());
                                emitter.send(SseEmitter.event()
                                        .name("error")
                                        .data(objectMapper.writeValueAsString(errorResponse)));
                            } catch (Exception sendEx) {
                            }
                            latch.countDown();
                        },
                        () -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                // 发送完成标记（包含最终token用量）
                                AnswerStreamResponse completeResponse = AnswerStreamResponse.content("[DONE]", true);
                                completeResponse.setPromptTokens(tokenUsage[0].get());
                                completeResponse.setCompletionTokens(tokenUsage[1].get());
                                completeResponse.setTotalTokens(tokenUsage[2].get());
                                emitter.send(SseEmitter.event()
                                        .name("completed")
                                        .data(objectMapper.writeValueAsString(completeResponse)));
                            } catch (Exception e) {
                            } finally {
                                latch.countDown();
                            }
                        });

                // 等待生成完成
                latch.await();

                emitter.complete();

            } catch (Exception e) {
                try {
                    AnswerStreamResponse errorResponse = AnswerStreamResponse.error("生成过程中发生错误: " + e.getMessage());
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data(objectMapper.writeValueAsString(errorResponse)));
                } catch (IOException sendEx) {
                    System.err.println("发送错误消息失败: " + sendEx.getMessage());
                }
                emitter.completeWithError(e);
            } finally {
                executor.shutdown();
            }
        });

        return emitter;
    }

    /**
     * 调用LLM生成回答
     */
    @Override
    public SseEmitter generateAnswerWithDocumentAndCollection(String topicId, String question,
            List<JSONObject> contexts) throws Exception {
        // 设置SSE超时时间为10分钟
        SseEmitter emitter = new SseEmitter(600000L);

        // 设置回调
        emitter.onCompletion(() -> {
        });
        emitter.onTimeout(() -> {
            emitter.complete();
        });
        emitter.onError(e -> System.err.println("SSE连接错误: " + e.getMessage()));

        // 使用独立线程池处理，避免阻塞主线程
        ExecutorService executor = Executors.newSingleThreadExecutor();

        executor.execute(() -> {
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                // 用于累计token用量
                AtomicLong[] tokenUsage = new AtomicLong[3];
                tokenUsage[0] = new AtomicLong(0); // promptTokens
                tokenUsage[1] = new AtomicLong(0); // completionTokens
                tokenUsage[2] = new AtomicLong(0); // totalTokens

                StringBuilder sb = new StringBuilder("招标文件内容：\n\n");
                CmcChat cmcChat = new CmcChat();
                cmcChat.setTopicId(topicId);
                for (CmcChat chat : cmcChatService.selectCmcChatList(cmcChat)) {
                    CmcDocument cmcDocument = new CmcDocument();
                    cmcDocument.setChatId(chat.getChatId());
                    List<CmcDocument> documentList = cmcDocumentService.selectCmcDocumentList(cmcDocument);
                    // 技术标书撰写智能体一个话题只会上传一个招标文件
                    if (documentList.size() == 1) {
                        for (CmcDocument document : documentList) {
                            File profilePath = new File(RuoYiConfig.getProfile() + document.getPath());
                            InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
                            List<TextSegment> segments = splitDocument(profilePath);
                            List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
                            embeddingStore.addAll(embeddings, segments);
                            Embedding queryEmbedding = embeddingModel.embed(question).content();
                            EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                                    .queryEmbedding(queryEmbedding)
                                    .minScore(0.7)
                                    .build();
                            List<EmbeddingMatch<TextSegment>> results = embeddingStore.search(embeddingSearchRequest)
                                    .matches();
                            results.sort(Comparator.comparingDouble(EmbeddingMatch<TextSegment>::score).reversed());
                            for (EmbeddingMatch<TextSegment> embeddingMatch : results) {
                                String requests = embeddingMatch.embedded().toString();
                                sb.append(requests).append("\n\n");
                            }
                        }
                    }
                }
                sb.append("针对本项目招标文件内容，补全以下章节部分：\n\n").append(question);

                String prompt = sb.toString();

                List<ChatMessage> messages = new ArrayList<>();
                if (topicId != null) {
                    CmcChat chat = new CmcChat();
                    chat.setTopicId(topicId);
                    List<CmcChat> cmcChatList = cmcChatService.selectCmcChatList(chat);
                    for (CmcChat c : cmcChatList) {
                        messages.add(ChatMessage.ofUser(c.getInput()));
                        messages.add(ChatMessage.ofAssistant(c.getOutput()));
                    }
                }
                messages.add(ChatMessage.ofUser(prompt));
                ChatSession chatSession = InMemoryChatSession.builder().messages(messages).build();
                Prompt prompt1 = Prompt.of(prompt).attrPut("session", chatSession);

                // 估算prompt的token数量（包含历史对话）
                StringBuilder fullPrompt = new StringBuilder();
                for (ChatMessage msg : messages) {
                    fullPrompt.append(msg.getContent());
                }
                tokenUsage[0].set((long) fullPrompt.length());

                // 使用 CountDownLatch 等待生成完成
                CountDownLatch latch = new CountDownLatch(1);
                AtomicBoolean cancelled = new AtomicBoolean(false);

                chatModel.prompt(prompt1).stream().map(response -> {
                    String content = response.lastChoice().getMessage().getContent();
                    // 累计completion tokens
                    if (content != null) {
                        tokenUsage[1].addAndGet((long) content.length());
                        tokenUsage[2].set(tokenUsage[0].get() + tokenUsage[1].get());
                    }
                    return content;
                }).subscribe(
                        content -> {
                            if (cancelled.get()) {
                                return;
                            }
                            try {
                                if (content != null && !content.isEmpty()) {
                                    AnswerStreamResponse response = AnswerStreamResponse.content(content, false);
                                    response.setPromptTokens(tokenUsage[0].get());
                                    response.setCompletionTokens(tokenUsage[1].get());
                                    response.setTotalTokens(tokenUsage[2].get());
                                    emitter.send(SseEmitter.event()
                                            .name("answer")
                                            .data(objectMapper.writeValueAsString(response)));
                                }
                            } catch (Exception e) {
                                cancelled.set(true);
                            }
                        },
                        error -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                AnswerStreamResponse errorResponse = AnswerStreamResponse
                                        .error("生成回答时发生错误: " + error.getMessage());
                                emitter.send(SseEmitter.event()
                                        .name("error")
                                        .data(objectMapper.writeValueAsString(errorResponse)));
                            } catch (Exception sendEx) {
                            }
                            latch.countDown();
                        },
                        () -> {
                            if (cancelled.getAndSet(true)) {
                                return;
                            }
                            try {
                                // 发送完成标记（包含最终token用量）
                                AnswerStreamResponse completeResponse = AnswerStreamResponse.content("[DONE]", true);
                                completeResponse.setPromptTokens(tokenUsage[0].get());
                                completeResponse.setCompletionTokens(tokenUsage[1].get());
                                completeResponse.setTotalTokens(tokenUsage[2].get());
                                emitter.send(SseEmitter.event()
                                        .name("completed")
                                        .data(objectMapper.writeValueAsString(completeResponse)));
                            } catch (Exception e) {
                            } finally {
                                latch.countDown();
                            }
                        });

                // 等待生成完成
                latch.await();

                emitter.complete();

            } catch (Exception e) {
                try {
                    AnswerStreamResponse errorResponse = AnswerStreamResponse.error("生成过程中发生错误: " + e.getMessage());
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data(objectMapper.writeValueAsString(errorResponse)));
                } catch (IOException sendEx) {
                    System.err.println("发送错误消息失败: " + sendEx.getMessage());
                }
                emitter.completeWithError(e);
            } finally {
                executor.shutdown();
            }
        });

        return emitter;
    }

    /**
     * 获取二级标题下三级标题列表
     */
    @Override
    public List<String> extractSubTitles(String filename, String question) throws IOException {
        List<String> subTitles = new ArrayList<>();
        boolean inTargetSection = false;

        InputStream fileInputStream = new FileInputStream(filename);
        try (XWPFDocument document = new XWPFDocument(fileInputStream)) {
            for (XWPFParagraph paragraph : document.getParagraphs()) {
                String text = paragraph.getText().trim();
                int level = getDocxOutlineLevel(paragraph, document);
                if (level == 3 && text.contains(question)) {
                    inTargetSection = true;
                    continue;
                }

                if (inTargetSection) {
                    if (level == 4) {
                        subTitles.add(text);
                    } else if (level == 3) {
                        break;
                    }
                }
            }
        }
        if (subTitles.size() == 0)
            subTitles.add(question);
        return subTitles;
    }

    /**
     * 检索知识库
     * 
     * @return
     */
    private List<List<SearchResp.SearchResult>> retrieve(String collectionName, String query, int topK) {
        MilvusClientV2 client = null;
        try {
            client = milvusConnectionPool.getClient();
            milvusConnectionPool.ensureCollectionLoaded(client, collectionName);

            List<BaseVector> queryVector = Collections
                    .singletonList(new FloatVec(embeddingModel.embed(query).content().vector()));

            Map<String, Object> searchParams = new HashMap<>();
            searchParams.put("nprobe", 8);

            SearchReqBuilder<?, ?> builder = SearchReq.builder()
                    .collectionName(collectionName)
                    .data(queryVector)
                    .topK(topK)
                    .outputFields(Arrays.asList("file_name", "file_type", "content"))
                    .annsField("embedding")
                    .metricType(IndexParam.MetricType.COSINE)
                    .searchParams(searchParams);

            SearchReq searchReq = builder.build();
            SearchResp searchResp = client.search(searchReq);
            return searchResp.getSearchResults();

        } catch (Exception e) {
            System.err.println("从Milvus检索时出错: " + e.getMessage());
            throw new RuntimeException("检索知识库失败", e);
        } finally {
            if (client != null) {
                milvusConnectionPool.returnClient(client);
            }
        }
    }

    /**
     * 提取文档中的所有二、三级标题
     */
    private String extractAllTitles(File transferFile) throws Exception {
        StringBuilder titlesBuilder = new StringBuilder();
        String filename = transferFile.getName().toLowerCase();

        if (filename.endsWith(".docx")) {
            try (XWPFDocument document = new XWPFDocument(new FileInputStream(transferFile))) {
                for (XWPFParagraph paragraph : document.getParagraphs()) {
                    String text = paragraph.getText().trim();
                    if (text.isEmpty())
                        continue;
                    int outlineLevel = getDocxOutlineLevel(paragraph, document);
                    if (outlineLevel == 2 || outlineLevel == 3) {
                        titlesBuilder.append(text).append("\n");
                    }
                }
            }
        } else if (filename.endsWith(".doc")) {
            try (HWPFDocument document = new HWPFDocument(new FileInputStream(transferFile))) {
                Range range = document.getRange();
                for (int i = 0; i < range.numParagraphs(); i++) {
                    Paragraph paragraph = range.getParagraph(i);
                    String text = paragraph.text().trim();
                    if (text.isEmpty())
                        continue;
                    int outlineLevel = getDocOutlineLevel(paragraph);
                    if (outlineLevel == 2 || outlineLevel == 3) {
                        titlesBuilder.append(text).append("\n");
                    }
                }
            }
        }
        return titlesBuilder.toString();
    }

    /**
     * 获取DOCX段落的大纲级别
     */
    private int getDocxOutlineLevel(XWPFParagraph paragraph, XWPFDocument document) {
        int level = 0;
        String styleId = paragraph.getStyleID();
        if (styleId != null) {
            XWPFStyle style = document.getStyles().getStyle(styleId);
            if (style != null) {
                CTPPrGeneral stylePPr = style.getCTStyle().getPPr();
                if (stylePPr != null && stylePPr.getOutlineLvl() != null) {
                    level = stylePPr.getOutlineLvl().getVal().intValue() + 1;
                }
            }
        }
        return level;
    }

    /**
     * 获取DOC段落的大纲级别
     */
    private int getDocOutlineLevel(Paragraph paragraph) {
        int level = 0;
        short styleIndex = paragraph.getStyleIndex();
        if (styleIndex >= 1 && styleIndex <= 9) {
            level = styleIndex;
        }
        return level;
    }

    /**
     * 检查DOCX段落是否是三级标题
     */
    private boolean isDocxLevel3Title(String paraText, XWPFDocument document) {
        for (XWPFParagraph p : document.getParagraphs()) {
            if (p.getText().trim().equals(paraText.trim())) {
                return getDocxOutlineLevel(p, document) == 3;
            }
        }
        return false;
    }

    /**
     * 按三级标题分割DOCX内容
     */
    private void splitDocxByLevel3(String content, XWPFDocument document, List<TextSegment> segments) {
        StringBuilder currentLevel3Content = new StringBuilder();
        String[] paragraphs = content.split("\n");
        for (String para : paragraphs) {
            if (para.trim().isEmpty())
                continue;

            if (isDocxLevel3Title(para.trim(), document)) {
                // 保存当前三级标题内容
                if (currentLevel3Content.length() != 0) {
                    TextSegment segment = TextSegment.from(currentLevel3Content.toString());
                    segments.add(segment);
                    currentLevel3Content = new StringBuilder();
                }
                // 开始新的三级标题内容
                currentLevel3Content.append(para).append("\n");
            } else {
                // 普通内容，添加到当前三级标题
                if (currentLevel3Content.length() != 0) {
                    currentLevel3Content.append(para).append("\n");
                }
            }
        }
        // 保存最后一个三级标题内容
        if (currentLevel3Content.length() != 0) {
            TextSegment segment = TextSegment.from(currentLevel3Content.toString());
            segments.add(segment);
        }
    }

    /**
     * 按三级标题分割DOC内容
     */
    private void splitDocByLevel3(String content, HWPFDocument document, List<TextSegment> segments) {
        StringBuilder currentLevel3Content = new StringBuilder();
        Range level2Range = document.getRange();
        boolean foundCurrentLevel2 = false;

        for (int j = 0; j < level2Range.numParagraphs(); j++) {
            Paragraph p = level2Range.getParagraph(j);
            String paraText = p.text().trim();
            if (paraText.isEmpty())
                continue;

            // 找到当前二级标题的开始
            int pLevel = getDocOutlineLevel(p);
            if (pLevel == 2 && paraText.equals(content.split("\n")[0].trim())) {
                foundCurrentLevel2 = true;
            }

            if (foundCurrentLevel2) {
                // 重新获取当前段落的大纲级别
                pLevel = getDocOutlineLevel(p);

                if (pLevel == 3) {
                    // 三级标题
                    // 保存当前三级标题内容
                    if (currentLevel3Content.length() != 0) {
                        TextSegment segment = TextSegment.from(currentLevel3Content.toString());
                        segments.add(segment);
                        currentLevel3Content = new StringBuilder();
                    }
                    // 开始新的三级标题内容
                    currentLevel3Content.append(paraText).append("\n");
                } else if (pLevel == 2 && !paraText.equals(content.split("\n")[0].trim())) {
                    // 下一个二级标题，结束当前二级标题的处理
                    break;
                } else {
                    // 普通内容或当前二级标题
                    if (currentLevel3Content.length() != 0 || pLevel == 2) {
                        currentLevel3Content.append(paraText).append("\n");
                    }
                }
            }
        }

        // 保存最后一个三级标题内容
        if (currentLevel3Content.length() != 0) {
            TextSegment segment = TextSegment.from(currentLevel3Content.toString());
            segments.add(segment);
        }
    }

    /**
     * 检索知识库
     */
    private List<TextSegment> splitDocument(File transferFile) throws Exception {
        // 加载文档
        Document document;
        InputStream fileInputStream = new FileInputStream(transferFile);
        String filename = transferFile.getName().toLowerCase();
        if (filename.endsWith(".docx")) {
            // 使用XWPFDocument处理DOCX文件，按段落层级分割
            List<TextSegment> segments = new ArrayList<>();
            try (XWPFDocument xwpfDocument = new XWPFDocument(fileInputStream)) {
                StringBuilder currentLevel2Content = new StringBuilder();
                boolean inLevel2Section = false;

                for (XWPFParagraph paragraph : xwpfDocument.getParagraphs()) {
                    String text = paragraph.getText().trim();
                    if (text.isEmpty())
                        continue;

                    int paraLevel = getDocxOutlineLevel(paragraph, xwpfDocument);

                    // 二级标题（大纲级别2）开始新的分段
                    if (paraLevel == 2) {
                        // 保存当前二级标题下的内容
                        if (currentLevel2Content.length() != 0) {
                            // 检查当前二级标题内容的字节数
                            int level2Length = currentLevel2Content.toString().getBytes().length;
                            if (level2Length <= 65535) {
                                // 字数不大于65535，按二级标题分割
                                TextSegment segment = TextSegment.from(currentLevel2Content.toString());
                                segments.add(segment);
                            } else {
                                // 字数大于65535，按三级标题分割
                                splitDocxByLevel3(currentLevel2Content.toString(), xwpfDocument, segments);
                            }
                            currentLevel2Content = new StringBuilder();
                        }
                        // 开始新的二级标题内容
                        currentLevel2Content.append(text).append("\n");
                        inLevel2Section = true;
                    }
                    // 其他层级（包括三级标题）
                    else {
                        if (inLevel2Section) {
                            currentLevel2Content.append(text).append("\n");
                        }
                    }
                }

                // 保存最后一个二级标题内容
                if (currentLevel2Content.length() != 0) {
                    // 检查字节数
                    int level2Length = currentLevel2Content.toString().getBytes().length;
                    if (level2Length <= 65535) {
                        // 字数不大于65535，按二级标题分割
                        TextSegment segment = TextSegment.from(currentLevel2Content.toString());
                        segments.add(segment);
                    } else {
                        // 字数大于65535，按三级标题分割
                        splitDocxByLevel3(currentLevel2Content.toString(), xwpfDocument, segments);
                    }
                }
            }
            return segments;
        } else if (filename.endsWith(".doc")) {
            // 使用HWPFDocument处理DOC文件，按段落层级分割
            List<TextSegment> segments = new ArrayList<>();
            try (HWPFDocument hwpfDocument = new HWPFDocument(fileInputStream)) {
                StringBuilder currentLevel2Content = new StringBuilder();
                boolean inLevel2Section = false;

                Range range = hwpfDocument.getRange();
                for (int i = 0; i < range.numParagraphs(); i++) {
                    Paragraph paragraph = range.getParagraph(i);
                    String text = paragraph.text().trim();
                    if (text.isEmpty())
                        continue;

                    int paraLevel = getDocOutlineLevel(paragraph);

                    if (paraLevel == 2) {
                        // 二级标题，开始新的分段
                        // 保存当前二级标题下的内容
                        if (currentLevel2Content.length() != 0) {
                            // 检查当前二级标题内容的字节数
                            int level2Length = currentLevel2Content.toString().getBytes().length;
                            if (level2Length <= 65535) {
                                // 字数不大于65535，按二级标题分割
                                TextSegment segment = TextSegment.from(currentLevel2Content.toString());
                                segments.add(segment);
                            } else {
                                // 字数大于65535，按三级标题分割
                                splitDocByLevel3(currentLevel2Content.toString(), hwpfDocument, segments);
                            }
                            currentLevel2Content = new StringBuilder();
                        }
                        // 开始新的二级标题内容
                        currentLevel2Content.append(text).append("\n");
                        inLevel2Section = true;
                    } else {
                        // 非二级标题或普通内容段落
                        if (inLevel2Section) {
                            currentLevel2Content.append(text).append("\n");
                        }
                    }
                }

                // 保存最后一个二级标题内容
                if (currentLevel2Content.length() != 0) {
                    // 检查字节数
                    int level2Length = currentLevel2Content.toString().getBytes().length;
                    if (level2Length <= 65535) {
                        // 字数不大于65535，按二级标题分割
                        TextSegment segment = TextSegment.from(currentLevel2Content.toString());
                        segments.add(segment);
                    } else {
                        // 字数大于65535，按三级标题分割
                        splitDocByLevel3(currentLevel2Content.toString(), hwpfDocument, segments);
                    }
                }
            }
            return segments;
        } else if (filename.endsWith(".pdf")) {
            document = new ApachePdfBoxDocumentParser().parse(fileInputStream);
        } else if (filename.endsWith(".txt")) {
            document = new TextDocumentParser().parse(fileInputStream);
        } else {
            throw new UnsupportedOperationException("不支持文件类型: " + filename);
        }
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(300, 50);
        return splitter.split(document);
    }

}