package com.ruoyi.web.llm.controller;

import com.alibaba.fastjson2.JSONObject;
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.McpChannel;
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 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());

    /**
     * 自动调用mcp工具问答
     * @return
     */
    @GetMapping("/answer")
    public AssistantMessage answer(String topicId, String question) throws IOException {
        McpClientProvider clientProvider = McpClientProvider.builder()
                .channel(McpChannel.SSE)
                .apiUrl("http://localhost:8080/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);
                arguments.put("title", question);
            }
            resultContent = clientProvider.callToolAsText(name, arguments).getContent();
            assistantMessage = new AssistantMessage(resultContent);
        }
        else
            throw new IOException("模型上下文工具未准备就绪，请重试");
        return assistantMessage;
    }


}
