package cn.com.poc.thirdparty.resource.demand.ai.function.chart_generate;

import cn.com.poc.agent_application.entity.KnowledgeContentResult;
import cn.com.poc.agent_application.entity.Variable;
import cn.com.poc.common.utils.JsonUtils;
import cn.com.poc.common.utils.StringUtils;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dbchain.DBChainResult;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dialogue.Message;
import cn.com.poc.thirdparty.resource.demand.ai.entity.largemodel.LargeModelDemandResult;
import cn.com.poc.thirdparty.resource.demand.ai.entity.largemodel.LargeModelResponse;
import cn.com.poc.thirdparty.resource.demand.ai.function.AbstractFunctionResult;
import cn.com.poc.thirdparty.resource.demand.ai.function.AbstractLargeModelFunction;
import cn.com.poc.thirdparty.resource.demand.ai.function.entity.FunctionLLMConfig;
import cn.com.poc.thirdparty.resource.demand.ai.function.entity.Parameters;
import cn.com.poc.thirdparty.resource.demand.ai.function.entity.Properties;
import cn.com.poc.thirdparty.service.LLMService;
import cn.hutool.core.collection.ListUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

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

/**
 * @author alex.yao
 * @date 2025/5/29
 */
@Component
public class ChartGenerateFunction extends AbstractLargeModelFunction {

    @Resource
    private LLMService llmService;
    private final Logger logger = LoggerFactory.getLogger(ChartGenerateFunction.class);

    private final String MODEL = "deepseek-v3";

    private final String DESC = "根据用户问题和现有数据生成mermaid图表,如生成饼状图,折线图,散点图,柱状图,流程图等";

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig.FunctionLLMConfigBuilder()
            .name("chart_generate")
            .description(DESC)
            .parameters(new Parameters("object")
                    .addProperties("question", new Properties("string", "用户问题优化后的问题"))
                    .addProperties("data", new Properties("string", "如果用户问题中存在数据则解析或截取对应的数据"))
                    .addProperties("theme", new Properties("string", "根据用户问题和上下文生成的图表主题"))
                    .addProperties("chart_type", new Properties("string", "图表类型", ListUtil.toList("饼状图", "折线图", "散点图", "柱状图", "流程图"))))

            .build();

    private final String SYSTEM_PROMPT = "您是一名专业的数据可视化工程师，请根据以下要求生成图表配置：\n" +
            "**主题**：\n" +
            "${theme}\n" +
            "**输入数据**：\n" +
            "上下文提取数据:\n" +
            "${data}\n" +
            "\n" +
            "数据库数据:\n" +
            "${database_result}\n" +
            "\n" +
            "知识库数据:\n" +
            "${knowledge_result}\n" +
            "\n" +
            "**图形类型**：\n" +
            "${chart_type}\n" +
            "\n" +
            "**配置要求**：\n" +
            "1. 图表标题应准确反映数据主题\n" +
            "2. 使用与数据类型匹配的图表子类型（如堆积柱状图、面积折线图等）\n" +
            "3. 坐标轴需要包含单位说明\n" +
            "4. 注意样式，如x轴内容过多时适当调整x轴宽度\n" +
            "5. 实际使用中用户的数据和要求不同，不能直接使用例子进行输出，要按照用户要求的图表进行\n" +
            "6. 当数据点超过5个时，x轴采用以下优化策略\n" +
            "        a) 使用config配置xyChart中的width调整适当宽度\n" +
            "        b) 调整xAxis的labelFontSize字体大小\n" +
            "注意不要照搬例子宽度！需要自行思考配置，例子如下：\n" +
            "---\n" +
            "config:\n" +
            "    xyChart:\n" +
            "        width: 1950\n" +
            "        height: 600\n" +
            "        xAxis:\n" +
            "            labelFontSize: 8\n" +
            "---\n" +
            "xychart-beta\n" +
            "......" +
            "\n" +
            "**输出要求**：\n" +
            "- 仅返回可直接用于Markdown的Mermaid配置\n" +
            "- 包含完整的Mermaid配置项\n" +
            "- 严格遵循Mermaid的语法\n" +
            "- 注意中文标题需要带双引号\n" +
            "- 若需要换行则使用</br>代替换行符\n" +
            "请以Mermaid格式返回配置, 思考好了之后再输出，不要输出其他内容，如解释或注释。\n" +
            "\n" +
            "**示例**\n" +
            "**示例1**\n" +
            "问题:帮我将一下数据生成折线图：\n" +
            "输入数据:请求类型 处理部门 平均处理时间 优先级 技术问题 技术部 2小时 高 账单问题 财务部 1工作日 中 其他 客服部 30分钟 低\n" +
            "主题：不同类型请求在各部门的平均处理时间对比\n" +
            "图形类型 ：柱状图\n" +
            "输出：\n" +
            "```mermaid\n" +
            "xychart-beta\n" +
            "    title \"不同类型请求在各部门的平均处理时间对比\"\n" +
            "    x-axis \"请求类型\" [\"技术问题\", \"账单问题\", \"其他\"]\n" +
            "    y-axis \"处理时间\" 0-->10\n" +
            "    bar [2, 8, 0.5]```\n" +
            "```\n" +
            "\n" +
            "**示例2**\n" +
            "问题: 帮我将一下数据生成折线图\n" +
            "输入数据: 2022-03-15 500000.00 定金 银行转账 YXSK20220315001 1 购房定金 2022-04-01 1500000.00 首付 银行转账 YXSK20220401001 1 首付款30% 2022-10-01 2450000.00 尾款 银行转账 YXSK20221001001 1 银行按揭放款 2022-04-20 500000.00 定金 银行转账 YXSK20220420001 2 购房定金 2022-05-05 1500000.00 首付 银行转账 YXSK20220505001 2 首付款30% 2022-11-01 2550000.00 尾款 银行转账 YXSK20221101001 2 银行按揭放款\n" +
            "主题：购房款项支付情况\n" +
            "图形类型 ：折线图\n" +
            "输出：\n" +
            "```mermaid\n" +
            "---\n" +
            "config:\n" +
            "    xyChart:\n" +
            "        width: 1950\n" +
            "        height: 600\n" +
            "        xAxis:\n" +
            "            labelFontSize: 11\n" +
            "---\n" +
            "xychart-beta\n" +
            "    title \"购房款项支付情况\"\n" +
            "    x-axis \"日期\" [\"2022-03-15\", \"2022-04-01\", \"2022-04-20\", \"2022-05-05\", \"2022-10-01\", \"2022-11-01\"]\n" +
            "    y-axis \"金额(元)\" 0-->3000000\n" +
            "    line [500000, 1500000, 500000, 1500000, 2450000, 2550000]\n" +
            "```";

    @Override
    public AbstractFunctionResult<String> doFunction(String content, String identifier, List<DBChainResult> dbChainResults, List<KnowledgeContentResult> knowledgeContentResults) {

        AbstractFunctionResult<String> result = new AbstractFunctionResult<>();
        result.setFunctionResult(StringUtils.EMPTY);
        result.setPromptContent(StringUtils.EMPTY);
        if (StringUtils.isBlank(content)) {
            return result;
        }
        JSONObject jsonObject = JSON.parseObject(content);
        if (!jsonObject.containsKey("question") || !jsonObject.containsKey("chart_type")) {
            logger.warn("缺少必要参数,无法生成图表,content:{}", content);
            return result;
        }

        String question = jsonObject.getString("question");

        Message systemMessage = new Message();
        systemMessage.setRole("system");
        systemMessage.setContent(buildPrompt(jsonObject, dbChainResults, knowledgeContentResults));

        Message userMessage = new Message();
        userMessage.setContent(question);
        userMessage.setRole("user");
        Message[] messages = new Message[]{systemMessage, userMessage};

        LargeModelResponse largeModelResponse = new LargeModelResponse();
        largeModelResponse.setModel(MODEL);
        largeModelResponse.setMessages(messages);
        largeModelResponse.setStream(false);
        largeModelResponse.setUser("Chart_Function");
        LargeModelDemandResult largeModelDemandResult = llmService.chat(largeModelResponse);
        if (largeModelDemandResult == null) {
            result.setFunctionResult(StringUtils.EMPTY);
            result.setPromptContent(StringUtils.EMPTY);
            return result;
        }
        result.setFunctionResult(largeModelDemandResult.getMessage());
        result.setPromptContent(largeModelDemandResult.getMessage());
        return result;
    }

    @Override
    public String getDesc() {
        return DESC;
    }

    @Override
    public List<String> getLLMConfig() {
        return ListUtil.toList(JsonUtils.serialize(functionLLMConfig));
    }

    @Override
    public List<String> getLLMConfig(List<Variable> variableStructure) {
        return this.getLLMConfig();
    }

    private String buildPrompt(JSONObject jsonObject, List<DBChainResult> dbChainResults, List<KnowledgeContentResult> knowledgeContentResults) {
        String chartType = StringUtils.EMPTY;
        String theme = StringUtils.EMPTY;
        String data = StringUtils.EMPTY;
        String sqlResult = StringUtils.EMPTY;
        String knowledgeResult = StringUtils.EMPTY;

        if (jsonObject.containsKey("chart_type")) {
            chartType = jsonObject.getString("chart_type");
        }
        if (jsonObject.containsKey("theme")) {
            theme = jsonObject.getString("theme");
        }

        if (jsonObject.containsKey("data")) {
            data = jsonObject.getString("data");
        }

        if (CollectionUtils.isNotEmpty(dbChainResults)) {
            sqlResult = dbChainResults.get(0).getSqlResult();
        }
        if (CollectionUtils.isNotEmpty(knowledgeContentResults)) {
            StringBuilder stringBuilder = new StringBuilder();
            int index = 1;
            for (KnowledgeContentResult knowledgeContentResult : knowledgeContentResults) {
                String content = knowledgeContentResult.getContent();
                stringBuilder.append("Chunk ").append(index++).append(":").append(content).append("\n\n");
            }
            knowledgeResult = stringBuilder.toString();
        }

        return SYSTEM_PROMPT.replace("${chart_type}", chartType)
                .replace("${theme}", theme)
                .replace("{data}", data)
                .replace("${database_result}", sqlResult)
                .replace("${knowledge_result}", knowledgeResult);
    }
}
