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

import cn.com.poc.agent_application.entity.KnowledgeContentResult;
import cn.com.poc.agent_application.entity.Variable;
import cn.com.poc.common.utils.DocumentLoad;
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.springframework.stereotype.Component;

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

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

    @Resource
    private LLMService llmService;

    private final String MODEL = "qwen-long";

    private final String DESC = "长文档理解，适合用于文档内容长、文件大的文档，结合用户问题与文档理解的插件,仅支持文档doc、docx、pdf、txt、md、xlsx、csv、xls";

    private final String SYSTEM_PROMPT = "# 工作规范：\n 1.工作流程：①对提供的文档内容进行理解，支持信息检索、摘要总结、文本分析。②根据用户提出的问题，提取或者总结文档中与问题相关的内容。2.工作限制：①要将问题与文档内容精准匹配，在理解文档时要带着问题去理解 \n\n";

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig
            .FunctionLLMConfigBuilder()
            .name("long_document_reader")
            .description(DESC)
            .parameters(new Parameters("object")
                    .addProperties("question", new Properties("string", "用户的问题"))
                    .addProperties("file_url", new Properties("string", "doc、docx、pdf、txt、md、xlsx、csv、xls文件地址"))

            ).build();

    @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("file_url")) {
            return result;
        }

        String question = jsonObject.getString("question");
        String fileUrl = jsonObject.getString("file_url");
        File file = DocumentLoad.downloadURLDocument(fileUrl);
        String documentContent;
        try {
            documentContent = DocumentLoad.documentToText(file);
        } catch (Exception e) {
            return result;
        }
        Message systemMessage = new Message();
        systemMessage.setRole("system");
        systemMessage.setContent(SYSTEM_PROMPT);

        Message documentMessage = new Message();
        documentMessage.setRole("system");
        documentMessage.setContent(documentContent);

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

        LargeModelResponse largeModelResponse = new LargeModelResponse();
        largeModelResponse.setModel(MODEL);
        largeModelResponse.setMessages(messages);
        largeModelResponse.setStream(false);
        largeModelResponse.setUser("Long_Document_Reader");
        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(this.functionLLMConfig));
    }

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