package com.trafficaudit.llmintegration.service;

import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.trafficaudit.llmintegration.config.DeepSeekConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
public class LLMService {

    private static final int MAX_TOOL_ROUNDS = 10;

    @Resource
    private DeepSeekConfig config;
    @Resource
    private DbQueryExecutor dbQueryExecutor;

    public String chat(String prompt) {
        JSONObject body = JSONUtil.createObj()
            .set("model", config.getModel())
            .set("messages", new JSONArray()
                .put(JSONUtil.createObj()
                    .set("role", "user")
                    .set("content", prompt)))
            .set("temperature", 0.3)
            .set("max_tokens", 1024);
        JSONObject result = post(body);
        if (result == null) {
            return null;
        }
        JSONArray choices = result.getJSONArray("choices");
        if (choices != null && !choices.isEmpty()) {
            return choices.getJSONObject(0).getJSONObject("message").getStr("content");
        }
        return null;
    }

    /**
     * 多轮对话：支持 query_database 工具调用，AI 可查询数据库后再回答。
     */
    public String chatWithTools(List<Map<String, Object>> messages, Map<String, Object> context) {
        JSONArray msgs = new JSONArray();
        JSONObject sys = new JSONObject();
        sys.set("role", "system");
        sys.set("content", buildSystemPrompt(context));
        msgs.put(sys);
        for (Map<String, Object> m : messages) {
            JSONObject j = new JSONObject();
            j.set("role", m.get("role"));
            j.set("content", m.get("content"));
            msgs.put(j);
        }

        JSONObject dbTool = buildQueryDatabaseTool();
        String period = context.get("reportPeriod") == null ? null : String.valueOf(context.get("reportPeriod"));

        for (int round = 0; round < MAX_TOOL_ROUNDS; round++) {
            JSONObject body = JSONUtil.createObj()
                .set("model", config.getModel())
                .set("messages", msgs)
                .set("tools", new JSONArray().put(dbTool))
                .set("tool_choice", "auto")
                .set("temperature", 0.3)
                .set("max_tokens", 2048);
            JSONObject result = post(body);
            if (result == null) {
                return "AI 服务调用失败，请稍后重试";
            }
            JSONArray choices = result.getJSONArray("choices");
            if (choices == null || choices.isEmpty()) {
                return "AI 服务返回异常，请稍后重试";
            }
            JSONObject message = choices.getJSONObject(0).getJSONObject("message");
            JSONArray toolCalls = message.getJSONArray("tool_calls");
            if (toolCalls == null || toolCalls.isEmpty()) {
                String content = message.getStr("content");
                return content == null ? "（AI 未返回内容）" : content;
            }

            // 回传 assistant 消息（含工具调用），再逐个执行并回填 tool 结果
            JSONObject assistant = new JSONObject();
            assistant.set("role", "assistant");
            assistant.set("content", message.getStr("content"));
            JSONArray calls = new JSONArray();
            for (int i = 0; i < toolCalls.size(); i++) {
                JSONObject tc = toolCalls.getJSONObject(i);
                JSONObject call = new JSONObject();
                call.set("id", tc.getStr("id"));
                call.set("type", "function");
                JSONObject fn = new JSONObject();
                fn.set("name", tc.getJSONObject("function").getStr("name"));
                fn.set("arguments", tc.getJSONObject("function").getStr("arguments"));
                call.set("function", fn);
                calls.put(call);
            }
            assistant.set("tool_calls", calls);
            msgs.put(assistant);

            for (int i = 0; i < toolCalls.size(); i++) {
                JSONObject tc = toolCalls.getJSONObject(i);
                String toolName = tc.getJSONObject("function").getStr("name");
                String args = tc.getJSONObject("function").getStr("arguments");
                log.info("LLM tool round {}: {} {}", round, toolName,
                    args != null && args.length() > 400 ? args.substring(0, 400) + "..." : args);
                String content;
                if ("query_database".equals(toolName)) {
                    JSONObject parsed;
                    try {
                        parsed = JSONUtil.parseObj(args);
                    } catch (Exception e) {
                        parsed = new JSONObject();
                    }
                    content = dbQueryExecutor.execute(
                        parsed.getStr("table"),
                        parsed.get("columns"),
                        parsed.get("filters"),
                        parsed.getStr("orderBy"),
                        parsed.getInt("limit"),
                        period);
                } else {
                    content = "{\"success\":false,\"error\":\"未知工具: " + toolName + "\"}";
                }
                JSONObject toolMsg = new JSONObject();
                toolMsg.set("role", "tool");
                toolMsg.set("tool_call_id", tc.getStr("id"));
                toolMsg.set("content", content);
                msgs.put(toolMsg);
            }
        }
        return "已达到最大查询轮次，请简化问题后重试";
    }

    private String buildSystemPrompt(Map<String, Object> ctx) {
        StringBuilder sb = new StringBuilder();
        sb.append("你是交通运输统计系统的 AI 助手，正在协助审核人员研判企业填报的核实解释（企业解释）是否合理、是否需要进一步核实。\n");
        if (ctx != null) {
            if (ctx.get("reportPeriod") != null) {
                sb.append("当前报表期: ").append(ctx.get("reportPeriod")).append("\n");
            }
            if (ctx.get("enterpriseName") != null) {
                sb.append("涉及企业: ").append(ctx.get("enterpriseName")).append("\n");
            }
            if (ctx.get("ruleDesc") != null && !ctx.get("ruleDesc").toString().trim().isEmpty()) {
                sb.append("关联审核规则: ").append(ctx.get("ruleDesc")).append("\n");
            }
            if (ctx.get("verifyExplanation") != null) {
                sb.append("企业解释: ").append(ctx.get("verifyExplanation")).append("\n");
            }
            if (ctx.get("reportId") != null) {
                sb.append("关联上报记录ID: ").append(ctx.get("reportId")).append("\n");
            }
        }
        sb.append("\n你可以调用 query_database 工具查询数据库获取真实数据，如企业车辆数、吨位、货运量、周转量、轨迹里程、规上规下拆分、去年同期等。\n\n");
        sb.append("数据库表结构如下（列名一律使用 camelCase，例如 enterpriseName、vehicleTotal、freightTotal、turnoverTotal；filters 的键同样使用 camelCase）：\n");
        sb.append(dbQueryExecutor.schemaText());
        sb.append("\n查询提示：");
        sb.append("1) 尽量先用 filters 精确过滤（如 enterpriseName 模糊匹配、reportPeriod 精确匹配、id 精确匹配），不要全表查询；");
        sb.append("2) 报表期使用 yyyy-MM 格式，例如 2026-07，若用户未指明报表期，优先使用当前报表期；");
        sb.append("3) 若首次查询报错，根据错误信息中给出的可用列名修正后重试，不要反复使用不存在的列名。\n\n");
        sb.append("基于查询到的数据给出有理有据的分析和处置建议，不要凭空猜测。回答请使用中文，并使用 Markdown 结构化排版，便于阅读：\n");
        sb.append("1) 用 # 或 ## 作为小节标题，例如「## 结论」「## 数据依据」「## 处置建议」；\n");
        sb.append("2) 关键结论用 **加粗** 强调，重要数字可直接加粗；\n");
        sb.append("3) 数据与要点用 - 项目符号或 1. 编号列表逐条列出，避免大段连续文字；\n");
        sb.append("4) 如需对比多组数据，可使用 Markdown 表格（| 表头 | 表头 |）。\n");
        sb.append("整体结构：先给结论，再列数据依据，最后给出处置建议。");
        return sb.toString();
    }

    private JSONObject buildQueryDatabaseTool() {
        JSONObject tableParam = JSONUtil.createObj()
            .set("type", "string")
            .set("enum", new JSONArray(Arrays.asList(DbQueryExecutor.ALLOWED_TABLES.toArray())))
            .set("description", "要查询的数据表，表名和列名见系统提示中的表结构");
        JSONObject parameters = JSONUtil.createObj()
            .set("type", "object")
            .set("properties", JSONUtil.createObj()
                .set("table", tableParam)
                .set("columns", JSONUtil.createObj()
                    .set("type", "array")
                    .set("items", JSONUtil.createObj().set("type", "string"))
                    .set("description", "要查询的列名（camelCase，必须来自系统提示中的表结构），不传则返回常用列"))
                .set("filters", JSONUtil.createObj()
                    .set("type", "object")
                    .set("description", "过滤条件，键为列名（camelCase），值为匹配值：字符串模糊匹配、数字精确匹配、数组为 IN 查询，也支持 '>=100' 这类比较写法"))
                .set("orderBy", JSONUtil.createObj()
                    .set("type", "string")
                    .set("description", "排序，如 id DESC"))
                .set("limit", JSONUtil.createObj()
                    .set("type", "integer")
                    .set("description", "返回条数上限，默认20，最大100")))
            .set("required", new JSONArray(Arrays.asList("table")));
        JSONObject fn = JSONUtil.createObj()
            .set("name", "query_database")
            .set("description", "查询交通统计数据数据库，获取真实数据用于分析。可用表: " + DbQueryExecutor.ALLOWED_TABLES)
            .set("parameters", parameters);
        return JSONUtil.createObj().set("type", "function").set("function", fn);
    }

    private JSONObject post(JSONObject body) {
        try {
            HttpResponse response = HttpRequest.post(config.getApiUrl())
                .header("Authorization", "Bearer " + config.getApiKey())
                .header("Content-Type", "application/json")
                .body(body.toString())
                .timeout(60000)
                .execute();
            if (response.isOk()) {
                return JSONUtil.parseObj(response.body());
            }
            log.error("LLM API error: {}", response.body());
            return null;
        } catch (Exception e) {
            log.error("LLM call failed", e);
            return null;
        }
    }
}