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

import com.alibaba.fastjson2.JSONObject;
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.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import io.milvus.client.MilvusClient;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResults;
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.InsertParam;
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.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.*;
import java.util.stream.Collectors;

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

    @Autowired
    private ICmcDocumentService cmcDocumentService;

    /**
     * 导入知识库文件
     */
    @Override
    public R<MutationResult> insertLangchainEmbeddingDocument(MilvusClient milvusClient, MultipartFile file, String collectionName, EmbeddingModel embeddingModel) 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<String> fileNames = new ArrayList<>();
        List<String> fileTypes = new ArrayList<>();
        List<String> texts = new ArrayList<>();
        List<List<Float>> embeddings = new ArrayList<>();

        for (TextSegment segment : segments) {
            String text = segment.text();
            if (text.trim().isEmpty())
                continue;
            fileNames.add(file.getOriginalFilename());
            String[] fileName = file.getOriginalFilename().split("\\.");
            fileTypes.add(fileName[fileName.length - 1]);
            texts.add(text);
            embeddings.add(embeddingModel.embed(text).content().vectorAsList());
        }

        // 准备导入数据
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("file_name", fileNames));
        fields.add(new InsertParam.Field("file_type", fileTypes));
        fields.add(new InsertParam.Field("content", texts));
        fields.add(new InsertParam.Field("embedding", embeddings));

        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(fields)
                .build();

        // 执行导入
        return milvusClient.insert(insertParam);
    }

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

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

    /**
     * 调用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(EmbeddingModel embeddingModel, 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(EmbeddingModel embeddingModel, 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;
    }

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

}