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.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.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;
import org.noear.solon.Solon;
import org.noear.solon.ai.annotation.PromptMapping;
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.mcp.server.annotation.McpServerEndpoint;
import org.noear.solon.annotation.Param;
import org.reactivestreams.Publisher;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

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 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 title,
                                                 @Param(description = "招标文件地址") String templatePath) throws IOException
    {
        title = String.join(",", extractSubTitles( "/upload/agent/template/technical.docx", title));
        List<JSONObject> requests = retrieveFromMilvus(milvusClient, embeddingModel, collectionName, title, 10);
        return generateAnswerWithDocumentAndCollection(embeddingModel, templatePath, title, requests, "http://192.168.28.188:8000/v1/chat/completions");
    }

    /**
     * 从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 templatePath, String question, List<JSONObject> contexts, String llmServiceUrl) throws IOException {
        StringBuilder sb = new StringBuilder("招标文件内容：\n\n");
        File profilePath = new File(Solon.cfg().getProperty("cmc.profile") + templatePath);
        List<TextSegment> segments = splitDocument(profilePath);
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings, segments);
        Embedding queryEmbedding = embeddingModel.embed(question).content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(3)
                .build();
        for (EmbeddingMatch embeddingMatch : embeddingStore.search(embeddingSearchRequest).matches()) {
            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(), llmServiceUrl);
    }

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

        chatSession.addMessage(ChatMessage.ofUser(prompt));

        ChatResponse response = chatModel.prompt(chatSession).call();
        return response.lastChoice().getMessage();
    }

    /**
     * 获取二级标题下三级标题列表
     */
    public List<String> extractSubTitles(String filename, String question) throws IOException {
        List<String> subTitles = new ArrayList<>();
        boolean inTargetSection = false;
        filename = Solon.cfg().getProperty("cmc.profile") + filename;
        InputStream fileInputStream = new FileInputStream(filename);
        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;
    }

    /**
     * 检索知识库
     */
    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(".docx")) {
            XWPFDocument docx = new XWPFDocument(fileInputStream);
            XWPFWordExtractor extractor = new XWPFWordExtractor(docx);
            String text = extractor.getText();
            document = Document.from(text);
        }
        else if (filename.endsWith(".pdf")) {
            document = new ApachePdfBoxDocumentParser().parse(fileInputStream);
        }
        else {
            throw new UnsupportedOperationException("不支持文件类型: " + filename);
        }
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(1000,200);
        return splitter.split(document);
    }

}