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.client.MilvusClient;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.collection.ReleaseCollectionParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.response.SearchResultsWrapper;
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.ChatSessionDefault;
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.server.annotation.McpServerEndpoint;
import org.noear.solon.annotation.Param;
import org.springframework.stereotype.Service;

import java.io.*;
import java.util.*;
import java.util.stream.Collectors;

@Service
@McpServerEndpoint(sseEndpoint = "/llm/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 MilvusServiceClient milvusClient = new MilvusServiceClient(
            ConnectParam.newBuilder()
                    .withHost("192.168.28.188")
                    .withPort(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"));
                title = String.join(",", extractSubTitles(templatePath, title));
//                List<JSONObject> contexts = retrieveFromMilvus(milvusClient, embeddingModel, collectionName, title, 10);
//                return generateAnswerWithDocumentAndCollection(embeddingModel, agentName, templatePath, title, contexts, llmServiceUrl);
                return generateAnswerWithDocumentAndCollection(embeddingModel, agentName, templatePath, title, new ArrayList<>(), llmServiceUrl);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
    }

    /**
     * 从Milvus检索相关文档
     * @return
     */
    public List<JSONObject> retrieveFromMilvus(MilvusClient milvusClient, EmbeddingModel embeddingModel, String collectionName, String query, int topK) {
        SearchResultsWrapper wrapper = retrieve(milvusClient, embeddingModel, collectionName, query, topK);
        return wrapper.getRowRecords(0).stream()
                .map(record -> {
                    JSONObject result = new JSONObject();
                    result.put("file_name", record.get("file_name"));
                    result.put("content", record.get("content"));
                    return result;
                })
                .collect(Collectors.toList());
    }

    /**
     * 调用LLM生成回答
     */
    public AssistantMessage generateAnswerWithDocumentAndCollection(EmbeddingModel embeddingModel, String agentName, String templatePath, String question, List<JSONObject> contexts, String llmServiceUrl) 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);
        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);
//        for (JSONObject context : contexts) {
//            sb.append("文件").append(": ")
//                    .append(context.getString("file_name")).append("\n\n")
//                    .append("段落格式").append(": ")
//                    .append(context.getString("content")).append("\n\n");
//        }
        return generateAnswer(sb.toString(), question, templatePath, llmServiceUrl);
    }

    /**
     * 调用LLM生成回答
     * @return
     */
    public AssistantMessage generateAnswer(String prompt, String question, String templatePath, String llmServiceUrl) 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();
        ChatResponse response = chatModel.prompt(chatSession).call();
        String content = response.lastChoice().getMessage().getContent() + "\n\n" +
                "招标文件分析完成，章节内容已写入【<a href='" + templatePath.replace(Solon.cfg().getProperty("cmc.profile"), "/dev-api/profile") + "'> 技术文件" + "</a>】，请查阅\n\n" +
                "如需修改，请输入技术文件已有内容的章节标题\n\n" +
                extractTitles(templatePath);
        String absolutePath = templatePath.replace("/dev-api/profile", Solon.cfg().getProperty("cmc.profile"));
        writeContent(response.lastChoice().getMessage().getContent(), question, absolutePath);
        return ChatMessage.ofAssistant(content);
    }

    /**
     * 写入章节内容
     * @return
     */
    public void writeContent(String content, String question, String absolutePath) throws IOException {
        String[] contentLines = content.split("\n");
        Map<String, String> map = new HashMap<>();
        String[] titles = question.split(",");
        File file = new File(absolutePath);
        FileInputStream fileInputStream = new FileInputStream(file);
        try (XWPFDocument document = new XWPFDocument(fileInputStream)) {
            for (int i = 0; i < titles.length; i++) {
                int startIndex = Arrays.asList(contentLines).indexOf(titles[i]);
                StringBuilder text = new StringBuilder();
                if (startIndex >= 0) {
                    if (i < titles.length - 1) {
                        int endIndex = Arrays.asList(contentLines).indexOf(titles[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[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<>();
        boolean inTargetSection = false;
        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") &&
                            text.contains(question)) {
                        inTargetSection = true;
                        continue;
                    }

                    // 如果已经在目标节中，收集标题3级别的子标题
                    if (inTargetSection) {
                        if (paragraph.getStyle().equals("4")) {
                            subTitles.add(text);
                        }
                        // 遇到下一个Heading1则退出
                        else if (paragraph.getStyle().equals("3")) {
                            break;
                        }
                    }
                }
            }
        }
        if (subTitles.size() == 0)
            subTitles.add(question);
        return subTitles;
    }

    /**
     * 获取二、三级标题列表
     */
    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();
    }

    /**
     * 检索知识库
     */
    private SearchResultsWrapper retrieve(MilvusClient milvusClient, EmbeddingModel embeddingModel, String collectionName, String query, int topK) {
        List<List<Float>> queryVector = Collections.singletonList(embeddingModel.embed(query).content().vectorAsList());

        //  加载集合
        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();

        R<RpcStatus> loadResponse = milvusClient.loadCollection(loadParam);
        if (loadResponse.getStatus() != R.Status.Success.getCode()) {
            System.err.println("加载Collection失败: " + loadResponse.getMessage());
            milvusClient.close();
        }

        // 构建SearchParam
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withVectors(queryVector)
                .withTopK(topK)
                .withOutFields(Arrays.asList("file_name", "file_type", "content"))
                .withVectorFieldName("embedding")
                .withMetricType(MetricType.COSINE)
                .withParams("{\"nprobe\": 8}")
                .build();

        R<SearchResults> response = milvusClient.search(searchParam);
        SearchResultsWrapper wrapper = new SearchResultsWrapper(response.getData().getResults());

        // 释放集合
        ReleaseCollectionParam param = ReleaseCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        milvusClient.releaseCollection(param);

        return wrapper;
    }

    /**
     * 检索知识库
     */
    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);
    }

}