package com.ruoyi.llm.service.impl;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.llm.domain.CmcChat;
import com.ruoyi.llm.service.ICmcChatService;
import com.ruoyi.llm.domain.AnswerStreamResponse;
import com.ruoyi.llm.service.ILangChainMilvusService;
import com.ruoyi.llm.service.ISessionService;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.noear.solon.ai.chat.ChatModel;
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.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

@Service
public class SessionServiceImpl implements ISessionService {

    @Autowired
    private ICmcChatService cmcChatService;

    @Autowired
    private ILangChainMilvusService langChainMilvusService;

    @Autowired
    private JdbcTemplate jdbcTemplate;

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

    /**
     * 获取所有表名列表（只包含cmc、sys和files开头的表）
     */
    private String getAllTableNames() {
        JSONArray tableNames = new JSONArray();
        try {
            String sql = "SELECT table_name, table_comment FROM information_schema.tables WHERE table_schema = DATABASE() "
                    +
                    "AND (table_name LIKE 'cmc%' OR table_name LIKE 'sys%' OR table_name LIKE 'files%')";
            List<Map<String, Object>> tables = jdbcTemplate.queryForList(sql);

            for (Map<String, Object> table : tables) {
                JSONObject tableInfo = new JSONObject();
                tableInfo.put("tableName", table.get("table_name"));
                tableInfo.put("tableComment", table.get("table_comment") != null ? table.get("table_comment") : "");
                tableNames.add(tableInfo);
            }
        } catch (Exception e) {
            throw new RuntimeException("获取表名列表失败", e);
        }
        return tableNames.toJSONString();
    }

    /**
     * 获取指定表的结构信息
     */
    private String getTableStructure(String tableName) {
        JSONObject result = new JSONObject();

        jdbcTemplate.execute((Connection conn) -> {
            try {
                DatabaseMetaData metaData = conn.getMetaData();

                ResultSet columns = metaData.getColumns(null, null, tableName, null);
                JSONArray columnsInfo = new JSONArray();

                while (columns.next()) {
                    JSONObject column = new JSONObject();
                    column.put("columnName", columns.getString("COLUMN_NAME"));
                    column.put("dataType", columns.getString("TYPE_NAME"));
                    column.put("columnSize", columns.getString("COLUMN_SIZE"));
                    column.put("isNullable", columns.getString("IS_NULLABLE"));
                    column.put("columnComment", columns.getString("REMARKS"));
                    columnsInfo.add(column);
                }
                columns.close();

                ResultSet primaryKeys = metaData.getPrimaryKeys(null, null, tableName);
                JSONArray primaryKeyColumns = new JSONArray();

                while (primaryKeys.next()) {
                    primaryKeyColumns.add(primaryKeys.getString("COLUMN_NAME"));
                }
                primaryKeys.close();

                result.put("tableName", tableName);
                result.put("columns", columnsInfo);
                result.put("primaryKey", primaryKeyColumns);

            } catch (SQLException e) {
                throw new RuntimeException("获取表结构失败: " + tableName, e);
            }
            return null;
        });

        return result.toJSONString();
    }

    /**
     * 执行SQL查询
     */
    private String executeSQL(String sqlString) {
        String cleanedSqlString = sqlString.replace("[", "").replace("]", "");

        JSONArray result = new JSONArray();

        jdbcTemplate.execute((Connection conn) -> {
            try (PreparedStatement stmt = conn.prepareStatement(cleanedSqlString)) {
                ResultSet rs = stmt.executeQuery();
                ResultSetMetaData metaData = rs.getMetaData();
                int columnCount = metaData.getColumnCount();

                while (rs.next()) {
                    JSONObject row = new JSONObject();
                    for (int i = 1; i <= columnCount; i++) {
                        String columnName = metaData.getColumnLabel(i);
                        Object value = rs.getObject(i);
                        row.put(columnName, value);
                    }
                    result.add(row);
                }
                rs.close();
            } catch (SQLException e) {
                throw new RuntimeException("SQL执行失败", e);
            }
            return null;
        });

        return result.toJSONString();
    }

    @Override
    public SseEmitter answer(String topicId, String question) {
        // 设置SSE超时时间为10分钟
        SseEmitter emitter = new SseEmitter(600000L);

        // 设置回调
        emitter.onCompletion(() -> {
        });
        emitter.onTimeout(() -> {
            emitter.complete();
        });
        emitter.onError(e -> System.err.println("SSE连接错误: " + e.getMessage()));

        // 使用独立线程池处理，避免阻塞主线程
        ExecutorService executor = Executors.newSingleThreadExecutor();

        executor.execute(() -> {
            ObjectMapper objectMapper = new ObjectMapper();
            try {
                // 用于累计token用量
                AtomicLong[] tokenUsage = new AtomicLong[3];
                tokenUsage[0] = new AtomicLong(0); // promptTokens
                tokenUsage[1] = new AtomicLong(0); // completionTokens
                tokenUsage[2] = new AtomicLong(0); // totalTokens
                System.out.println("llmServiceUrl: " + llmServiceUrl);

                ChatModel chatModel = ChatModel.of(llmServiceUrl).model("Qwen")
                        .build();
                // ChatModel chatModel =
                // ChatModel.of("https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions")
                // .model("qwen3-vl-32b-instruct")
                // .apiKey("sk-750a17cc723847f28b31fa1bc17c255c")
                // .build();

                System.out.println("成功");
                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 获取所有表名（使用本地方法替代MCP）
                String toolResult1 = getAllTableNames();
                System.out.println(toolResult1);

                // 第二步：模型分析表名，找到相关表
                String step2Prompt = "请分析以下表名列表，找出与" + question + "相关的表。\n\n" +
                        "**输出格式要求**：\n" +
                        "- 如果需要数据库查询，请以固定格式输出：以\"相关表：\" 开头，中间用英文逗号隔开，以英文句号结尾\n" +
                        "- 如果不需要进行数据库查询，请直接回答。\n\n" + toolResult1;

                // 使用 CountDownLatch 等待生成完成
                CountDownLatch latch = new CountDownLatch(1);
                AtomicBoolean cancelled = new AtomicBoolean(false);
                StringBuilder step2ContentBuilder = new StringBuilder();

                // 生成第二步的响应
                Prompt step3Prompt = Prompt.of(step2Prompt).attrPut("session",
                        InMemoryChatSession.builder().messages(messages).build());
                chatModel.prompt(step3Prompt)
                        .stream()
                        .subscribe(
                                resp1 -> {
                                    if (cancelled.get()) {
                                        return;
                                    }
                                    // 收集流式响应内容
                                    String chunk = resp1.getContent();
                                    if (chunk != null && !chunk.isEmpty()) {
                                        step2ContentBuilder.append(chunk);
                                    }
                                },
                                error -> {
                                    if (cancelled.getAndSet(true)) {
                                        return;
                                    }
                                    try {
                                        emitter.send(SseEmitter.event()
                                                .name("error")
                                                .data("生成回答时发生错误: " + error.getMessage()));
                                    } catch (Exception sendEx) {
                                    }
                                    latch.countDown();
                                },
                                () -> {
                                    if (cancelled.get()) {
                                        return;
                                    }
                                    try {
                                        String content2 = step2ContentBuilder.toString();
                                        System.out.println(content2);
                                        // 提取相关表名（从完整的模型响应中提取）
                                        String relevantTables = extractRelevantTables(content2);
                                        System.out.println(relevantTables);

                                        // 检查是否需要继续执行（如果包含工具调用或明确需要数据库查询）
                                        if (relevantTables.equals("")) {
                                            // 不需要数据库查询，直接返回最终回答
                                            // 估算token用量
                                            StringBuilder fullPrompt = new StringBuilder();
                                            for (ChatMessage msg : messages) {
                                                fullPrompt.append(msg.getContent());
                                            }
                                            fullPrompt.append(step2Prompt);
                                            tokenUsage[0].set((long) fullPrompt.length());
                                            tokenUsage[1].set((long) content2.length());
                                            tokenUsage[2].set(tokenUsage[0].get() + tokenUsage[1].get());

                                            AnswerStreamResponse response = AnswerStreamResponse.content(content2,
                                                    false);
                                            response.setPromptTokens(tokenUsage[0].get());
                                            response.setCompletionTokens(tokenUsage[1].get());
                                            response.setTotalTokens(tokenUsage[2].get());
                                            emitter.send(SseEmitter.event()
                                                    .name("answer")
                                                    .data(objectMapper.writeValueAsString(response)));

                                            AnswerStreamResponse completeResponse = AnswerStreamResponse
                                                    .content("[DONE]", true);
                                            completeResponse.setPromptTokens(tokenUsage[0].get());
                                            completeResponse.setCompletionTokens(tokenUsage[1].get());
                                            completeResponse.setTotalTokens(tokenUsage[2].get());
                                            emitter.send(SseEmitter.event()
                                                    .name("completed")
                                                    .data(objectMapper.writeValueAsString(completeResponse)));
                                            latch.countDown();
                                            return;
                                        }

                                        // 第三步：调用 getTableStructure 获取相关表的结构（使用本地方法替代MCP）
                                        // 执行第三步：获取表结构
                                        StringBuilder tableStructures = new StringBuilder();
                                        for (String tableName : relevantTables.split(",")) {
                                            tableName = tableName.trim();
                                            if (!tableName.isEmpty()) {
                                                String toolResult3 = getTableStructure(tableName);
                                                tableStructures.append("表：").append(tableName).append("\n")
                                                        .append(toolResult3)
                                                        .append("\n\n");
                                            }
                                        }

                                        // 第四步：模型生成 SQL 查询
                                        String step4Prompt = "请根据以下表结构，生成" + question + "的SQL语句：\n\n"
                                                + tableStructures.toString();

                                        // 生成第四步的响应 - 同样需要收集完整响应后再提取SQL
                                        StringBuilder step4ContentBuilder = new StringBuilder();

                                        Prompt step5Prompt = Prompt.of(step4Prompt).attrPut("session",
                                                InMemoryChatSession.builder().messages(messages).build());
                                        chatModel.prompt(step5Prompt)
                                                .stream()
                                                .subscribe(
                                                        resp2 -> {
                                                            if (cancelled.get()) {
                                                                return;
                                                            }
                                                            // 收集流式响应内容
                                                            String chunk = resp2.getContent();
                                                            if (chunk != null && !chunk.isEmpty()) {
                                                                step4ContentBuilder.append(chunk);
                                                            }
                                                        },
                                                        step4Error -> {
                                                            if (cancelled.getAndSet(true)) {
                                                                return;
                                                            }
                                                            try {
                                                                emitter.send(SseEmitter.event()
                                                                        .name("error")
                                                                        .data("生成SQL时发生错误: "
                                                                                + step4Error.getMessage()));
                                                            } catch (Exception sendEx) {
                                                            }
                                                            latch.countDown();
                                                        },
                                                        () -> {
                                                            if (cancelled.get()) {
                                                                return;
                                                            }
                                                            try {
                                                                String content4 = step4ContentBuilder.toString();

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

                                                                // 第五步：调用 executeSQL 执行查询（使用本地方法替代MCP）
                                                                // 执行第五步：执行 SQL 查询
                                                                String toolResult5 = executeSQL(sqlQuery);

                                                                // 第六步：模型生成最终回答
                                                                String step6Prompt = "请根据以下查询结果，生成关于" + question
                                                                        + "的最终回答：\n\n"
                                                                        + toolResult5;

                                                                // 估算prompt的token数量（包含历史对话）
                                                                StringBuilder fullPrompt = new StringBuilder();
                                                                for (ChatMessage msg : messages) {
                                                                    fullPrompt.append(msg.getContent());
                                                                }
                                                                fullPrompt.append(step6Prompt);
                                                                tokenUsage[0].set((long) fullPrompt.length());

                                                                // 生成第六步的响应（流式）
                                                                Prompt step7Prompt = Prompt.of(step6Prompt).attrPut(
                                                                        "session",
                                                                        InMemoryChatSession.builder().messages(messages)
                                                                                .build());
                                                                chatModel.prompt(step7Prompt)
                                                                        .stream()
                                                                        .map(resp -> {
                                                                            String content = resp.getContent();
                                                                            // 累计completion tokens
                                                                            if (content != null) {
                                                                                tokenUsage[1].addAndGet(
                                                                                        (long) content.length());
                                                                                tokenUsage[2].set(tokenUsage[0].get()
                                                                                        + tokenUsage[1].get());
                                                                            }
                                                                            return content;
                                                                        })
                                                                        .subscribe(
                                                                                content -> {
                                                                                    if (cancelled.get()) {
                                                                                        return;
                                                                                    }
                                                                                    try {
                                                                                        if (content != null
                                                                                                && !content.isEmpty()) {
                                                                                            AnswerStreamResponse response = AnswerStreamResponse
                                                                                                    .content(content,
                                                                                                            false);
                                                                                            response.setPromptTokens(
                                                                                                    tokenUsage[0]
                                                                                                            .get());
                                                                                            response.setCompletionTokens(
                                                                                                    tokenUsage[1]
                                                                                                            .get());
                                                                                            response.setTotalTokens(
                                                                                                    tokenUsage[2]
                                                                                                            .get());
                                                                                            emitter.send(SseEmitter
                                                                                                    .event()
                                                                                                    .name("answer")
                                                                                                    .data(objectMapper
                                                                                                            .writeValueAsString(
                                                                                                                    response)));
                                                                                        }
                                                                                    } catch (Exception e) {
                                                                                        cancelled.set(true);
                                                                                    }
                                                                                },
                                                                                error -> {
                                                                                    if (cancelled.getAndSet(true)) {
                                                                                        return;
                                                                                    }
                                                                                    try {
                                                                                        AnswerStreamResponse errorResponse = AnswerStreamResponse
                                                                                                .error("生成回答时发生错误: "
                                                                                                        + error.getMessage());
                                                                                        emitter.send(SseEmitter.event()
                                                                                                .name("error")
                                                                                                .data(objectMapper
                                                                                                        .writeValueAsString(
                                                                                                                errorResponse)));
                                                                                    } catch (Exception sendEx) {
                                                                                    }
                                                                                    latch.countDown();
                                                                                },
                                                                                () -> {
                                                                                    if (cancelled.getAndSet(true)) {
                                                                                        return;
                                                                                    }
                                                                                    try {
                                                                                        // 发送完成标记（包含最终token用量）
                                                                                        AnswerStreamResponse completeResponse = AnswerStreamResponse
                                                                                                .content("[DONE]",
                                                                                                        true);
                                                                                        completeResponse
                                                                                                .setPromptTokens(
                                                                                                        tokenUsage[0]
                                                                                                                .get());
                                                                                        completeResponse
                                                                                                .setCompletionTokens(
                                                                                                        tokenUsage[1]
                                                                                                                .get());
                                                                                        completeResponse.setTotalTokens(
                                                                                                tokenUsage[2].get());
                                                                                        emitter.send(SseEmitter.event()
                                                                                                .name("completed")
                                                                                                .data(objectMapper
                                                                                                        .writeValueAsString(
                                                                                                                completeResponse)));
                                                                                    } catch (Exception e) {
                                                                                    }
                                                                                    latch.countDown();
                                                                                });
                                                            } catch (Exception e) {
                                                                if (cancelled.getAndSet(true)) {
                                                                    return;
                                                                }
                                                                try {
                                                                    emitter.send(SseEmitter.event()
                                                                            .name("error")
                                                                            .data("抱歉，系统执行查询时遇到问题，请稍后重试。"));
                                                                } catch (Exception sendEx) {
                                                                }
                                                                latch.countDown();
                                                            }
                                                        });
                                    } catch (Exception e) {
                                        if (cancelled.getAndSet(true)) {
                                            return;
                                        }
                                        try {
                                            emitter.send(SseEmitter.event()
                                                    .name("error")
                                                    .data("抱歉，系统执行查询时遇到问题，请稍后重试。"));
                                        } catch (Exception sendEx) {
                                        }
                                        latch.countDown();
                                    }
                                });

                // 等待生成完成
                latch.await();

                emitter.complete();

            } catch (Exception e) {
                try {
                    emitter.send(SseEmitter.event()
                            .name("error")
                            .data("生成过程中发生错误: " + e.getMessage()));
                } catch (IOException sendEx) {
                    System.err.println("发送错误消息失败: " + sendEx.getMessage());
                }
                emitter.completeWithError(e);
            } finally {
                executor.shutdown();
            }
        });

        return emitter;
    }

    /**
     * 从模型响应中提取相关表名
     */
    private String extractRelevantTables(String content) {
        if (content == null || content.trim().isEmpty()) {
            return "";
        }

        // 尝试匹配 "相关表：" 或 "相关表\n" 格式
        String[] markers = { "相关表：", "相关表", "相关表名：", "相关表名" };
        int start = -1;
        String matchedMarker = "";

        for (String marker : markers) {
            int idx = content.indexOf(marker);
            if (idx != -1 && (start == -1 || idx < start)) {
                start = idx;
                matchedMarker = marker;
            }
        }

        if (start != -1) {
            start += matchedMarker.length();
            // 跳过可能的换行符
            while (start < content.length() && (content.charAt(start) == '\n' || content.charAt(start) == '\r'
                    || content.charAt(start) == ' ' || content.charAt(start) == '：')) {
                start++;
            }

            // 找到表名结束位置（换行、句号或字符串结尾）
            int end = start;
            while (end < content.length() && content.charAt(end) != '\n' && content.charAt(end) != '\r'
                    && content.charAt(end) != '.') {
                end++;
            }

            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 SseEmitter answerWithDocument(String topicId, String chatId, String question) throws Exception {
        return langChainMilvusService.generateAnswerWithDocument(topicId, chatId, question);
    }

}