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

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.constants.LLMRoleEnum;
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.entity.largemodel.Thinking;
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.chart_generate.ChartGenerateFunction;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

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

/**
 * @author 52747
 * @date 2025/8/19
 */
@Component
public class EChartGenerateFunction extends AbstractLargeModelFunction {
    @Resource
    private LLMService llmService;

    private final Logger logger = LoggerFactory.getLogger(ChartGenerateFunction.class);

    private final String DESC = "根据用户问题、SQL和SQL执行结果,生成折线图,饼图,柱状图";

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig.FunctionLLMConfigBuilder()
            .name("eChart_generate")
            .description(DESC)
            .parameters(new Parameters("object")
                    .addProperties("question", new Properties("string", "用户问题优化后的问题"))
                    .addProperties("sql", new Properties("string", "如果用户问题中存在数据则解析或截取对应的数据"))
                    .addProperties("sql_result", new Properties("string", "根据用户问题和上下文生成的图表主题")))
            .build();

    private static String PROMPT = "## 角色\n" +
            "你的任务是根据数据和用户提问生成EChart 的Option Json\n" +
            "\n" +
            "## 限制\n" +
            "- 输出格式: 适用JSON格式输出\n" +
            "- 仅输出JSON内容\n" +
            "- 图表限制: 折线图、饼图、柱状图\n" +
            "- 如果数据与问题不能生成图表，则输出 : { \"skip\" : true }\n" +
            "\n" +
            "## 用户提问\n" +
            "${QUESTION}\n" +
            "\n" +
            "## SQL\n" +
            "${SQL}\n" +
            "\n" +
            "## SQL Result\n" +
            "${SQL_RESULT}";

    @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("sql") || !jsonObject.containsKey("sql_result")) {
            logger.warn("缺少必要参数,无法生成图表,content:{}", content);
            return result;
        }

        String question = jsonObject.getString("question");
        String sql = jsonObject.getString("sql");
        String sqlResult = jsonObject.getString("sql_result");
        String option = generateEChart(sql, sqlResult, question);
        JSONObject optionJSONObject = JSON.parseObject(option);
        if (optionJSONObject.containsKey("skip")) {
            return result;
        }

        result.setPromptContent(option);
        result.setFunctionResult(option);
        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 generateEChart(String SQL, String SQLResult, String question) {
        String model = "ep-20250814152748-6jlnx";
        String prompt = PROMPT.replace("${QUESTION}", question)
                .replace("${SQL}", SQL)
                .replace("${SQL_RESULT}", SQLResult);

        Message systemMessage = new Message();
        systemMessage.setRole(LLMRoleEnum.SYSTEM.getRole());
        systemMessage.setContent(prompt);

        Message userMessage = new Message();
        userMessage.setRole(LLMRoleEnum.USER.getRole());
        userMessage.setContent(question);


        LargeModelResponse largeModelResponse = new LargeModelResponse();
        largeModelResponse.setThinking(new Thinking() {{
            setType("disabled");
        }});
        largeModelResponse.setModel(model);
        largeModelResponse.setMessages(new Message[]{systemMessage, userMessage});
        largeModelResponse.setStream(false);
        largeModelResponse.setUser("GENERATE_E_CHART");
        LargeModelDemandResult largeModelDemandResult = llmService.chat(largeModelResponse);
        return largeModelDemandResult.getMessage();
    }
}
