package com.ruoyi.agent.service.impl;

import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.agent.service.IMcpService;
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.ConnectConfig;
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.SearchReq;
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.SearchResp;
import org.apache.poi.extractor.POITextExtractor;
import org.apache.poi.extractor.ExtractorFactory;
import org.apache.poi.xwpf.usermodel.*;
import org.apache.xmlbeans.XmlCursor;
import org.noear.solon.Solon;
import org.noear.solon.ai.annotation.ToolMapping;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatResponse;
import org.noear.solon.ai.chat.ChatSession;
import org.noear.solon.ai.chat.message.AssistantMessage;
import org.noear.solon.ai.chat.message.ChatMessage;
import org.noear.solon.ai.chat.session.InMemoryChatSession;
import org.noear.solon.ai.mcp.McpChannel;
import org.noear.solon.ai.mcp.server.annotation.McpServerEndpoint;
import org.noear.solon.annotation.Param;
import org.springframework.stereotype.Service;

import java.io.*;
import java.util.*;

@Service
@McpServerEndpoint(channel = McpChannel.SSE, mcpEndpoint = "/mcp/sse")
public class McpServiceImpl implements IMcpService {

    private static final EmbeddingModel embeddingModel = new BgeSmallZhV15EmbeddingModel();

    private static final String llmServiceUrl = "http://192.168.28.188:8000/v1/chat/completions";

    private static final MilvusClientV2 milvusClient = new MilvusClientV2(
            ConnectConfig.builder()
                    .uri("http://192.168.28.188:19530")
                    .build());

    /**
     * 调用LLM+RAG（外部文件+知识库）生成回答
     */
    @ToolMapping(description = "章节撰写")
    public AssistantMessage writeParagraph(@Param(description = "知识库名称") String collectionName,
                                           @Param(description = "智能体名称") String agentName,
                                           @Param(description = "章节名称") String title,
                                           @Param(description = "技术文件地址") String templatePath) throws IOException
    {
            try {
                templatePath = templatePath.replace("/dev-api/profile", Solon.cfg().getProperty("cmc.profile"));
                List<String> subTitles = extractSubTitles(templatePath, title);
                List<JSONObject> contexts = retrieveFromMilvus(collectionName, title, 10);
                return generateAnswerWithDocumentAndCollection(agentName, templatePath, subTitles, contexts);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
    }

    /**
     * 从Milvus检索相关文档
     * @return
     */
    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;
    }

    /**
     * 调用LLM生成回答
     */
    public AssistantMessage generateAnswerWithDocumentAndCollection(String agentName, String templatePath, List<String> titles, List<JSONObject> contexts) throws IOException {
        StringBuilder sb = new StringBuilder("招标文件内容：\n\n");
        String filename = templatePath.replace("_" + agentName, "");
        File profilePath = new File(filename);
        if (!profilePath.exists()) {
            filename = filename.replace(".docx", ".doc");
            profilePath = new File(filename);
        }
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        List<TextSegment> segments = splitDocument(profilePath);
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        embeddingStore.addAll(embeddings, segments);
        StringBuilder content = new StringBuilder();
        for (String title : titles) {
            Embedding queryEmbedding = embeddingModel.embed(title).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(title);
            content.append(generateAnswer(sb.toString()));
        }
        String absolutePath = templatePath.replace("/dev-api/profile", Solon.cfg().getProperty("cmc.profile"));
        writeContent(content.toString(), titles, absolutePath);
        for (JSONObject context : contexts) {
            sb.append("文件").append(": ")
                    .append(context.getString("file_name")).append("\n\n")
                    .append("段落格式").append(": ")
                    .append(context.getString("content")).append("\n\n");
        }
        content.append( "招标文件分析完成，章节内容已写入【<a href='")
                .append(templatePath.replace(Solon.cfg().getProperty("cmc.profile"), "/dev-api/profile"))
                .append("'> 技术文件" + "</a>】，请查阅\n\n")
                .append("如需修改，请输入技术文件已有内容的章节标题\n\n")
                .append(extractTitles(templatePath));
        return ChatMessage.ofAssistant(content.toString());
    }

    /**
     * 调用LLM生成回答
     * @return
     */
    public String generateAnswer(String prompt) throws IOException {
        ChatModel chatModel = ChatModel.of(llmServiceUrl)
                .model("Qwen2.5-1.5B-Instruct")
                .build();

        List<ChatMessage> messages = new ArrayList<>();
        messages.add(ChatMessage.ofUser(prompt));
        ChatSession chatSession =  InMemoryChatSession.builder().messages(messages).build();
        chatSession.addMessage(messages);
        ChatResponse response = chatModel.prompt(chatSession).call();

        return response.lastChoice().getMessage().getContent();
    }

    /**
     * 写入章节内容
     * @return
     */
    public void writeContent(String content, List<String> titles, String absolutePath) throws IOException {
        String[] contentLines = content.split("\n");
        Map<String, String> map = new HashMap<>();
        File file = new File(absolutePath);
        FileInputStream fileInputStream = new FileInputStream(file);
        try (XWPFDocument document = new XWPFDocument(fileInputStream)) {
            for (int i = 0; i < titles.size(); i++) {
                int startIndex = Arrays.asList(contentLines).indexOf(titles.get(i));
                StringBuilder text = new StringBuilder();
                if (startIndex >= 0) {
                    if (i < titles.size() - 1) {
                        int endIndex = Arrays.asList(contentLines).indexOf(titles.get(i + 1));
                        for (int c = startIndex + 1; c < endIndex; c++) {
                            text.append(contentLines[c]).append("\n\n");
                        }
                    } else {
                        if (startIndex + 1 < contentLines.length) {
                            for (int c = startIndex + 1; c < contentLines.length; c++) {
                                text.append(contentLines[c]).append("\n\n");
                            }
                        }
                    }
                }
                else
                    text.append(content);
                map.put(titles.get(i), text.toString());
            }

            List<Integer> positions = new ArrayList<>();
            List<String> contents = new ArrayList<>();
            List<XWPFParagraph> paragraphs = document.getParagraphs();

            for (int i = 0; i < paragraphs.size(); i++) {
                XWPFParagraph paragraph = paragraphs.get(i);
                for (String title : titles) {
                    if (paragraph.getText().contains(title)) {
                        positions.add(i);
                        contents.add(map.get(title));
                    }
                }
            }

            for (int i = positions.size() - 1; i >= 0; i--) {
                int insertPos = positions.get(i) + 1;
                XmlCursor xmlCursor = paragraphs.get(insertPos).getCTP().newCursor();
                XWPFParagraph contentParagraph = document.insertNewParagraph(xmlCursor);
                contentParagraph.setStyle("1");
                XWPFRun run = contentParagraph.createRun();
                run.setText(contents.get(i));
            }

            try (FileOutputStream out = new FileOutputStream(absolutePath)) {
                document.write(out);
            }
        }
    }

    /**
     * 获取最低级别子标题列表
     */
    public List<String> extractSubTitles(String filename, String question) throws IOException {
        List<String> subTitles = new ArrayList<>();
        InputStream fileInputStream = new FileInputStream(filename);

        boolean foundParent = false;
        int parentLevel = -1;
        // 用于跟踪当前路径
        List<XWPFParagraph> currentPath = new ArrayList<>();

        try (XWPFDocument document = new XWPFDocument(fileInputStream)) {

            for (XWPFParagraph  paragraph : document.getParagraphs()) {
                String text = paragraph.getText();
                int level = Integer.parseInt(paragraph.getStyle());

                // 维护当前路径
                while (!currentPath.isEmpty() && Integer.parseInt(currentPath.get(currentPath.size() - 1).getStyle()) >= level) {
                    currentPath.remove(currentPath.size() - 1);
                }
                currentPath.add(paragraph);
                if (foundParent) {
                    if (level > parentLevel) { // 是子标题
                        if (isLeafNode(document, paragraph, level)) {
                            subTitles.add(text);
                        }
                    } else {
                        // 遇到同级或更高级别的标题，停止搜索
                        break;
                    }
                } else if (text.equals(question)) {
                    foundParent = true;
                    parentLevel = level;
                }
            }


            return subTitles;

        }
    }

    // 检查一个标题是否是叶子节点（没有更低级别的子标题）
    private boolean isLeafNode(XWPFDocument doc, XWPFParagraph p, int level) {
        int index = doc.getPosOfParagraph(p);
        if (index == -1 || index >= doc.getParagraphs().size()-1) {
            return true;
        }

        // 检查后续段落
        XWPFParagraph nextP = doc.getParagraphs().get(index + 1);
        int nextLevel = Integer.parseInt(nextP.getStyle());
        // 遇到同级或更高级别标题，是叶子节点
        return nextLevel <= level; // 存在更低级别的标题，不是叶子节点

    }

    /**
     * 获取二、三级标题列表
     */
    public String extractTitles(String filename) throws IOException {
        StringBuilder subTitles = new StringBuilder();
        InputStream fileInputStream = new FileInputStream(filename);
        try (XWPFDocument document = new XWPFDocument(fileInputStream)) {
            for (XWPFParagraph paragraph : document.getParagraphs()) {
                String text = paragraph.getText().trim();
                if (paragraph.getStyle() != null) {
                    // 判断主标题
                    if (paragraph.getStyle().equals("3") || paragraph.getStyle().equals("4") ) {
                        subTitles.append(text).append("\n");
                    }
                }
            }
        }
        return subTitles.toString();
    }

    /**
     * 检索知识库
     * @return
     */
    private List<List<SearchResp.SearchResult>> retrieve(String collectionName, String query, int topK) {
        List<BaseVector> queryVector = Collections.singletonList(new FloatVec(embeddingModel.embed(query).content().vector()));

        //  加载集合
        LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                .collectionName(collectionName)
                .build();
        milvusClient.loadCollection(loadCollectionReq);

        // 构建SearchParam
        Map<String, Object> searchParams = new HashMap<>();
        searchParams.put("nprobe", 8);
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(queryVector)
                .topK(topK)
                .outputFields(Arrays.asList("file_name", "file_type", "content"))
                .annsField("embedding")
                .metricType(IndexParam.MetricType.COSINE)
                .searchParams(searchParams)
                .build();

        SearchResp searchResp = milvusClient.search(searchReq);
        List<List<SearchResp.SearchResult>> searchResultList = searchResp.getSearchResults();

        // 释放集合
        ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                .collectionName(collectionName)
                .build();
        milvusClient.releaseCollection(releaseCollectionReq);

        return searchResultList;
    }

    /**
     * 检索知识库
     */
    private List<TextSegment> splitDocument(File transferFile) throws IOException {
        // 加载文档
        Document document;
        InputStream fileInputStream = new FileInputStream(transferFile);
        String filename = transferFile.getName().toLowerCase();
        if (filename.endsWith(".doc") || filename.endsWith(".docx")) {
            try (POITextExtractor extractor = ExtractorFactory.createExtractor(fileInputStream)) {
                String text = extractor.getText();
                document = Document.from(text);
            }
            catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
        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);
    }

}