package com.ruoyi.web.llm.service.impl;

import com.alibaba.fastjson2.JSONObject;
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.llm.service.ICmcChatService;
import com.ruoyi.llm.service.ICmcDocumentService;
import com.ruoyi.web.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.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.InsertReq;
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.InsertResp;
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.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;
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.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

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

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

    @Autowired
    private ICmcDocumentService cmcDocumentService;

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

    private static final EmbeddingModel embeddingModel = new BgeSmallZhV15EmbeddingModel();

    /**
     * 导入知识库文件
     * @return
     */
    @Override
    public long insertLangchainEmbeddingDocument(MultipartFile file, String collectionName) throws IOException
    {
        File profilePath = new File( RuoYiConfig.getProfile() + "/upload/rag/knowledge/" + collectionName);
        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<>();

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

            JSONObject fastjsonObj = new JSONObject();
            fastjsonObj.put("file_name", file.getOriginalFilename());
            String[] fileName = file.getOriginalFilename().split("\\.");
            fastjsonObj.put("file_type", fileName[fileName.length - 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);
        }

        InsertReq insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(data)
                .build();

        // 执行导入
        InsertResp insertResp = milvusClient.insert(insertReq);
        return insertResp.getInsertCnt();
    }

    /**
     * 从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 Flux<AssistantMessage> generateAnswer(String topicId, String prompt, String llmServiceUrl) {
        ChatModel chatModel = ChatModel.of(llmServiceUrl)
                
                .model("Qwen2.5-1.5B-Instruct")
                .build();

        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();
        Publisher<ChatResponse> publisher = chatModel.prompt(chatSession).stream();
        return Flux.from(publisher)
                .map(response -> {
                    return response.lastChoice().getMessage();
                });
    }

    /**
     * 调用LLM+RAG生成回答
     * @return
     */
    @Override
    public Flux<AssistantMessage> generateAnswerWithCollection(String topicId, String question, List<JSONObject> contexts, String llmServiceUrl) {
        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");
        }
        return generateAnswer(topicId, sb.toString(), llmServiceUrl);
    }

    /**
     * 调用LLM生成回答
     */
    @Override
    public Flux<AssistantMessage> generateAnswerWithDocument(String topicId, String chatId, String question, String llmServiceUrl) throws IOException {
        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");
            }
        }
        return generateAnswer(topicId, sb.toString(), llmServiceUrl);
    }

    /**
     * 调用LLM生成回答
     */
    @Override
    public Flux<AssistantMessage> generateAnswerWithDocumentAndCollection(String topicId, String question, List<JSONObject> contexts, String llmServiceUrl) throws IOException {
        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);
//        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(topicId, sb.toString(), llmServiceUrl);
    }

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

    /**
     * 检索知识库
     * @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 FileNotFoundException {
        // 加载文档
        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);
    }

}