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

import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.llm.domain.CmcChat;
import com.ruoyi.llm.service.ICmcChatService;
import com.ruoyi.web.llm.service.ILangChainMilvusService;
import com.ruoyi.web.llm.service.ISessionService;
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.ChatMessage;
import org.noear.solon.ai.chat.prompt.Prompt;
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.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

@Service
public class SessionServiceImpl implements ISessionService {

    @Autowired
    private ICmcChatService cmcChatService;

    @Autowired
    private ILangChainMilvusService langChainMilvusService;

    @Value("${cmc.llmService.url}")
    private String llmServiceUrl;

    @Value("${cmc.mcpService.url}")
    private String mcpServiceUrl;

    @Override
    public Flux<String> answer(String topicId, String question) {
        McpClientProvider clientProvider = McpClientProvider.builder()
                .channel(McpChannel.STREAMABLE_STATELESS)
                .url(mcpServiceUrl)
                .build();
        ChatModel chatModel = ChatModel.of(llmServiceUrl)
                .model("Qwen")
                .defaultToolAdd(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()));
        }

        // 第一步：调用 GetAllTableNames 获取所有表名
        String step1Prompt = "你是一个MySQL专家。请调用 GetAllTableNames 工具获取数据库中所有表名列表。";
        messages.add(ChatMessage.ofUser(step1Prompt));
        ChatSession session1 = InMemoryChatSession.builder().messages(messages).build();
        Prompt prompt1 = Prompt.of(step1Prompt).attrPut("session", session1);

        Flux<ChatResponse> chatResponse = chatModel.prompt(prompt1).stream();
        Flux<String> contentFlux = chatResponse
                .concatMap(resp -> {
                    try {
                        // 执行第一步：获取所有表名
                        String toolCall1 = "{\"name\": \"GetAllTableNames\", \"arguments\": {}}";
                        String toolResult1 = clientProvider.callTool("GetAllTableNames", new JSONObject()).getContent();

                        // 第二步：模型分析表名，找到相关表
                        List<ChatMessage> messages2 = new ArrayList<>(messages);
                        messages2.add(ChatMessage.ofAssistant("<tool_call>" + toolCall1 + "</tool_call>"));
                        messages2.add(ChatMessage.ofUser(toolResult1));

                        String step2Prompt = "请分析以下表名列表，找出与" + question + "相关的表。\n\n" +
                        "**输出格式要求**：\n" +
                        "- 如果需要数据库查询，请以固定格式输出：以\"相关表：\" 开头，中间用英文逗号隔开，以英文句号结尾\n" +
                        "- 如果不需要进行数据库查询，请直接回答。\n\n" + toolResult1;
                        messages2.add(ChatMessage.ofUser(step2Prompt));
                        ChatSession session2 = InMemoryChatSession.builder().messages(messages2).build();

                        // 生成第二步的响应
                        return chatModel.prompt(Prompt.of(step2Prompt).attrPut("session", session2))
                                .stream()
                                .concatMap(resp1 -> {
                                    try {
                                        String content2 = resp1.getContent();

                                        // 提取相关表名（从模型响应中提取）
                                        String relevantTables = extractRelevantTables(content2);

                                        // 检查是否需要继续执行（如果包含工具调用或明确需要数据库查询）
                                        if (relevantTables.equals("")) {
                                            // 不需要数据库查询，直接返回最终回答
                                            return Flux.just(content2)
                                                    .map(this::toSseDataFrame)
                                                    .concatWith(Flux.just("data: [DONE]\n\n"));
                                        }

                                        // 第三步：调用 GetTableStructure 获取相关表的结构
                                        List<ChatMessage> messages3 = new ArrayList<>(messages2);
                                        messages3.add(ChatMessage.ofAssistant(content2));

                                        // 执行第三步：获取表结构
                                        StringBuilder tableStructures = new StringBuilder();
                                        for (String tableName : relevantTables.split(",")) {
                                            tableName = tableName.trim();
                                            if (!tableName.isEmpty()) {
                                                JSONObject args = new JSONObject();
                                                args.put("tableName", tableName);
                                                String toolResult3 = clientProvider.callTool("GetTableStructure", args)
                                                        .getContent();
                                                tableStructures.append("表：").append(tableName).append("\n")
                                                        .append(toolResult3)
                                                        .append("\n\n");
                                            }
                                        }

                                        // 第四步：模型生成 SQL 查询
                                        List<ChatMessage> messages4 = new ArrayList<>(messages3);
                                        messages4.add(ChatMessage.ofAssistant(
                                                "<tool_call>{\"name\": \"GetTableStructure\", \"arguments\": {\"tableName\": \"employee\"}}</tool_call>"));
                                        messages4.add(ChatMessage.ofUser(tableStructures.toString()));

                                        String step4Prompt = "请根据以下表结构，生成" + question + "的SQL语句：\n\n"
                                                + tableStructures.toString();
                                        messages4.add(ChatMessage.ofUser(step4Prompt));

                                        // 生成第四步的响应
                                        return chatModel
                                                .prompt(Prompt.of(step4Prompt).attrPut("session",
                                                        InMemoryChatSession.builder().messages(messages4).build()))
                                                .stream()
                                                .concatMap(resp2 -> {
                                                    try {
                                                        String content4 = resp2.getContent();

                                                        // 提取 SQL 语句
                                                        String sqlQuery = extractSqlFromContent(content4);

                                                        // 第五步：调用 SQLQuery 执行查询
                                                        List<ChatMessage> messages5 = new ArrayList<>(messages4);
                                                        messages5.add(ChatMessage.ofAssistant(content4));

                                                        // 执行第五步：执行 SQL 查询
                                                        JSONObject sqlArgs = new JSONObject();
                                                        sqlArgs.put("sqlString", sqlQuery);
                                                        String toolResult5 = clientProvider
                                                                .callTool("SQLQuery", sqlArgs)
                                                                .getContent();

                                                        // 第六步：模型生成最终回答
                                                        List<ChatMessage> messages6 = new ArrayList<>(messages5);
                                                        messages6.add(ChatMessage.ofAssistant(
                                                                "<tool_call>{\"name\": \"SQLQuery\", \"arguments\": {\"sqlString\": \""
                                                                        + sqlQuery + "\"}}</tool_call>"));
                                                        messages6.add(ChatMessage.ofUser(toolResult5));

                                                        String step6Prompt = "请根据以下查询结果，生成关于" + question + "的最终回答：\n\n"
                                                                + toolResult5;
                                                        messages6.add(ChatMessage.ofUser(step6Prompt));

                                                        // 生成第六步的响应（流式）
                                                        return chatModel
                                                                .prompt(Prompt.of(step6Prompt).attrPut("session",
                                                                        InMemoryChatSession.builder()
                                                                                .messages(messages6)
                                                                                .build()))
                                                                .stream()
                                                                .map(resp3 -> resp3.getContent())
                                                                .filter(answer -> answer != null && !answer.isEmpty());
                                                    } catch (Exception e) {
                                                        return Flux.just("抱歉，系统执行查询时遇到问题，请稍后重试。")
                                                                .map(this::toSseDataFrame)
                                                                .concatWith(Flux.just("data: [DONE]\n\n"));
                                                    }
                                                });
                                    } catch (Exception e) {
                                        return Flux.just("抱歉，系统执行查询时遇到问题，请稍后重试。")
                                                .map(this::toSseDataFrame)
                                                .concatWith(Flux.just("data: [DONE]\n\n"));
                                    }
                                });
                    } catch (Exception e) {
                        return Flux.just("抱歉，系统执行查询时遇到问题，请稍后重试。")
                                .map(this::toSseDataFrame)
                                .concatWith(Flux.just("data: [DONE]\n\n"));
                    }
                });
        // 手工拼接 SSE 帧，绕过 ServerSentEvent 编码差异；并在末尾发送 [DONE]
        Flux<String> sseFlux = contentFlux
                .map(this::toSseDataFrame)
                .concatWith(Flux.just("data: [DONE]\n\n"));
        return sseFlux;
    }

    /**
     * 从模型响应中提取相关表名
     */
    private String extractRelevantTables(String content) {
        // 简化处理，实际应该根据模型输出格式提取
        // 假设模型会返回类似 "相关表：employee, department"
        if (content.contains("相关表：")) {
            int start = content.indexOf("相关表：") + 4;
            int end = content.indexOf(".", start);
            if (end > start) {
                String tables = content.substring(start, end).trim();
                // 限制相关表最多只保留10个
                String[] tableNames = tables.split(",");
                if (tableNames.length > 10) {
                    StringBuilder limitedTables = new StringBuilder();
                    for (int i = 0; i < 10; i++) {
                        if (i > 0) {
                            limitedTables.append(",");
                        }
                        limitedTables.append(tableNames[i].trim());
                    }
                    return limitedTables.toString();
                }
                return tables;
            }
        }
        // 默认返回可能相关的表
        return "";
    }

    /**
     * 从模型响应中提取 SQL 语句
     */
    private String extractSqlFromContent(String content) {
        // 简化处理，实际应该根据模型输出格式提取
        if (content.contains("SELECT")) {
            int start = content.indexOf("SELECT");
            int end = content.indexOf(";", start);
            if (end > start) {
                return content.substring(start, end + 1);
            }
        }
        return "";
    }

    @Override
    public Flux<String> answerWithDocument(String topicId, String chatId, String question) throws Exception {
        return langChainMilvusService.generateAnswerWithDocument(topicId, chatId, question);
    }

    private String toSseDataFrame(String data) {
        if (data == null) {
            return "data: \n\n";
        }
        // SSE 要求每一行都以 data: 开头
        String normalized = data.replace("\r\n", "\n");
        StringBuilder sb = new StringBuilder(normalized.length() + 16);
        sb.append("data: ");
        int start = 0;
        while (true) {
            int idx = normalized.indexOf('\n', start);
            if (idx < 0) {
                sb.append(normalized.substring(start));
                break;
            }
            sb.append(normalized, start, idx);
            sb.append("\n");
            sb.append("data: ");
            start = idx + 1;
        }
        sb.append("\n\n");
        return sb.toString();
    }

}
