package com.ruoyi.web.llm.controller;

import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.common.config.RuoYiConfig;
import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.llm.domain.CmcChat;
import com.ruoyi.llm.service.ICmcAgentService;
import com.ruoyi.llm.service.ICmcChatService;
import com.ruoyi.llm.service.ICmcTopicService;
import com.ruoyi.web.llm.service.ILangChainMilvusService;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallzhv15.BgeSmallZhV15EmbeddingModel;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
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.client.McpClientProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

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


/**
 * mcp模型上下文协议Controller
 * 
 * @author cmc
 * @date 2025-04-08
 */
@RestController
@RequestMapping("/llm/mcp")
public class McpController extends BaseController
{
    @Autowired
    private ILangChainMilvusService langChainMilvusService;

    @Autowired
    private ICmcChatService cmcChatService;

    @Autowired
    private ICmcTopicService cmcTopicService;

    @Autowired
    private ICmcAgentService cmcAgentService;

    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());
//
//    /**
//     * 同步问答
//     * @return
//     */
//    @GetMapping("/answer")
//    public AssistantMessage answer(String question) throws IOException {
//        McpClientProvider clientProvider = McpClientProvider.builder()
//                .apiUrl("http://localhost:8080/llm/mcp/sse")
//                .build();
//        ChatModel chatModel = ChatModel.of(llmServiceUrl)
//                .model("Qwen2.5-1.5B-Instruct")
//                .defaultToolsAdd(clientProvider)
//                .build();
//        ChatResponse response = chatModel.prompt(question).call();
//        String resultContent = response.lastChoice().getMessage().getResultContent();
//        AssistantMessage assistantMessage = new AssistantMessage(resultContent);
//        if (resultContent.startsWith("<tool_call>")) {
//            String content = resultContent.replace("<tool_call>\n", "").replace("\n</tool_call>", "");
//            JSONObject jsonObject = JSONObject.parseObject(content);
//            String name = jsonObject.getString("name");
//            JSONObject arguments = jsonObject.getJSONObject("arguments");
//            resultContent = clientProvider.callToolAsText(name, arguments).getContent();
//            assistantMessage = new AssistantMessage(resultContent);
//        }
//        return assistantMessage;
//    }

    /**
     * 自动调用mcp工具问答
     * @return
     */
    @GetMapping("/answer")
    public AssistantMessage answer(String topicId, String question) throws IOException {
        McpClientProvider clientProvider = McpClientProvider.builder()
                .apiUrl("http://localhost:8080/llm/mcp/sse")
                .build();
        ChatModel chatModel = ChatModel.of(llmServiceUrl)
                .model("Qwen2.5-1.5B-Instruct")
                .defaultToolsAdd(clientProvider)
                .build();

        List<ChatMessage> messages = new ArrayList<>();
        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(question));
        ChatSession chatSession =  InMemoryChatSession.builder().messages(messages).build();
        ChatResponse response = chatModel.prompt(chatSession).call();
        String resultContent = response.lastChoice().getMessage().getResultContent();
        AssistantMessage assistantMessage;
        if (resultContent.startsWith("<tool_call>")) {
            String content = resultContent.replace("<tool_call>\n", "").replace("\n</tool_call>", "");
            JSONObject jsonObject = JSONObject.parseObject(content);
            String name = jsonObject.getString("name");
            JSONObject arguments = jsonObject.getJSONObject("arguments");
            if (arguments.getString("templatePath").contains("招标") || arguments.getString("templatePath").contains("询价")) {
                arguments.put("collectionName", "technical");
                String agentName = cmcAgentService.selectCmcAgentByAgentId(cmcTopicService.selectCmcTopicByTopicId(topicId).getAgentId()).getAgentName();
                arguments.put("agentName", agentName);
            }
            resultContent = clientProvider.callToolAsText(name, arguments).getContent();
            assistantMessage = new AssistantMessage(resultContent);
        }
        else
            throw new IOException("模型上下文工具未准备就绪，请重试");
        return assistantMessage;
    }

    /**
     * 调用LLM+RAG（外部文件+知识库）生成回答
     */
    @GetMapping("/answerWithDocument")
    public Flux<AssistantMessage> answerWithDocumentAndCollection(String collectionName, String topicId, String question) throws IOException
    {
        question = String.join(",", langChainMilvusService.extractSubTitles(RuoYiConfig.getProfile() + "/upload/agent/template/technical.docx", question));
        List<JSONObject> requests = langChainMilvusService.retrieveFromMilvus(milvusClient, embeddingModel, collectionName, question, 10);
        return langChainMilvusService.generateAnswerWithDocumentAndCollection(embeddingModel, topicId, question, requests, llmServiceUrl);
    }

    /**
     * 根据招标文件要求编写投标文件技术方案
     * @return
     */
    @GetMapping("/tech")
    public AssistantMessage tech(MultipartFile file) throws IOException {
        McpClientProvider clientProvider = McpClientProvider.builder()
                    .apiUrl("http://localhost:8080/llm/mcp/sse")
                    .build();
        ChatModel chatModel = ChatModel.of(llmServiceUrl)
                
                .model("DeepSeek-R1-Distill-Qwen-1.5B")
                .defaultToolsAdd(clientProvider)
                .build();
        Map<String,Object> path = new HashMap<>();
        String content = clientProvider.callToolAsText("openDocument", path).getContent();
        Map<String,Object> request = new HashMap<>();
        request.put("request", content);
        List<ChatMessage> messages = clientProvider.getPromptAsMessages("askQuestion", request);
        ChatResponse response = chatModel.prompt(messages).call();
        System.out.println(response.getChoices().get(0).getMessage());
        Map<String,Object> write = new HashMap<>();
        write.put("resultContent", response.getChoices().get(0).getMessage().getResultContent());
        String res = clientProvider.callToolAsText("writeTechnicalPlan", write).getContent();
        return response.getChoices().get(0).getMessage();
    }

}
