package cn.com.poc.thirdparty.resource.demand.ai.function.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.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.hutool.core.collection.ListUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.springframework.stereotype.Component;

import java.io.File;
import java.util.List;

@Component
public class DocumentReaderFunction extends AbstractLargeModelFunction {

    private final String DESC = "文档阅读器，读取PDF、docx、doc、txt、md格式文件";

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig
            .FunctionLLMConfigBuilder()
            .name("document_reader")
            .description(DESC)
            .parameters(new Parameters("object")
                    .addProperties("file_urls", new Properties("array", "doc、docx、pdf、txt、md文件地址"))

            ).build();


    @Override
    public AbstractFunctionResult<String> doFunction(String content, String identifier, List<DBChainResult> dbChainResults, List<KnowledgeContentResult> knowledgeContentResults) {
        AbstractFunctionResult<String> result = new AbstractFunctionResult<>();
        if (StringUtils.isBlank(content)) {
            result.setFunctionResult(StringUtils.EMPTY);
            result.setPromptContent(StringUtils.EMPTY);
            return result;
        }
        StringBuilder sb = new StringBuilder();
        JSONObject jsonObject = JSON.parseObject(content);
        if (jsonObject.containsKey("file_urls")) {
            JSONArray jsonArray = jsonObject.getJSONArray("file_urls");
            int size = jsonArray.size();
            if (size == 0) {
                result.setFunctionResult(StringUtils.EMPTY);
                result.setPromptContent(StringUtils.EMPTY);
                return result;
            }
            for (int i = 0; i < size; i++) {
                String fileUrl = jsonArray.getString(i);
                File file = DocumentLoad.downloadURLDocument(fileUrl);
                sb.append(StringUtils.LF).append("## Document ").append((i + 1)).append(StringUtils.LF);
                sb.append(DocumentLoad.documentToText(file));
            }
            result.setFunctionResult(sb.toString());
            result.setPromptContent(sb.toString());
            return result;
        }
        result.setFunctionResult(StringUtils.EMPTY);
        result.setPromptContent(StringUtils.EMPTY);
        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();
    }
}
