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

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.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.com.poc.thirdparty.resource.demand.ai.function.extraction.entity.KeyInfo;
import cn.com.poc.thirdparty.resource.textin.api.TextInClient;
import cn.hutool.core.collection.ListUtil;
import cn.hutool.json.JSONException;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

/**
 * 合同关键信息抽取-要素提取
 *
 * @author alex.yao
 * @date 2025/5/12
 */
@Component
public class ContractExtractionFunction extends AbstractLargeModelFunction {

    private final String DESC = "合同关键信息抽取";

    private final TextInClient textInClient = new TextInClient();

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig.FunctionLLMConfigBuilder()
            .name("contract_extraction")
            .parameters(new Parameters("array")
                    .addProperties("file_url", new Properties("string", "文件链接, 合同文件的在线地址"))
                    .addProperties("key_info", new Properties("string", "关键信息名称, 长度限制20个字符"))
                    .addProperties("paraphrase_names", new Properties("array", "相似名字段,字符串数组, 可根据相似名精准抽取关键信息, 最多填写3个,每个释义名称长度限制20个字符"))
                    .addProperties("field_type", new Properties("string", "字段类型字段, 可选项有,时间:time, 金额:amount, 地址:address, 公司:company, 姓名:name, 描述(长文本):long_text_description, 其他:other, 印章:stamp, 分别对应产品段配置的字段类型"))
                    .addProperties("keywords", new Properties("array", "关键字字段, 字符串数组, 可根据关键字信息，快速定位抽取信所在段落范围, 最多填写10个,且字符总长度不超过50"))
            )
            .description(DESC)
            .build();

    @Override
    public AbstractFunctionResult<String> doFunction(String content, String identifier, List<DBChainResult> dbChainResults, List<KnowledgeContentResult> knowledgeContentResults) {
        AbstractFunctionResult<String> result = new AbstractFunctionResult<>();
        String fileUrl;
        List<KeyInfo> keyInfos = new ArrayList<>();
        if (isJsonArray(content)) {
            JSONArray jsonArray = JSONArray.parseArray(content);
            if (jsonArray.isEmpty()) {
                return result;
            }
            fileUrl = jsonArray.getJSONObject(0).getString("file_url");
            for (int i = 0; i < jsonArray.size(); i++) {
                JSONObject jsonObject = jsonArray.getJSONObject(i);
                KeyInfo keyInfo = new KeyInfo();
                if (jsonObject.containsKey("field_type")) {
                    keyInfo.setField_type(jsonObject.getString("file_type"));
                }
                if (jsonObject.containsKey("key_info")) {
                    keyInfo.setKey_info(jsonObject.getString("key_info"));
                }
                if (jsonObject.containsKey("paraphrase_names")) {
                    keyInfo.setParaphrase_names(jsonObject.getJSONArray("paraphrase_names").toArray(new String[0]));
                }
                if (jsonObject.containsKey("keywords")) {
                    keyInfo.setKeywords(jsonObject.getJSONArray("keywords").toArray(new String[0]));
                }
                keyInfos.add(keyInfo);
            }
        } else {
            JSONObject jsonObject = JSONObject.parseObject(content);
            fileUrl = jsonObject.getString("file_url");
            KeyInfo keyInfo = new KeyInfo();
            if (jsonObject.containsKey("field_type")) {
                keyInfo.setField_type(jsonObject.getString("file_type"));
            }
            if (jsonObject.containsKey("key_info")) {
                keyInfo.setKey_info(jsonObject.getString("key_info"));
            }
            if (jsonObject.containsKey("paraphrase_names")) {
                keyInfo.setParaphrase_names(jsonObject.getJSONArray("paraphrase_names").toArray(new String[0]));
            }
            if (jsonObject.containsKey("keywords")) {
                keyInfo.setKeywords(jsonObject.getJSONArray("keywords").toArray(new String[0]));
            }
            keyInfos.add(keyInfo);
        }
        String extraction = textInClient.extraction(fileUrl, keyInfos);
        result.setFunctionResult(extraction);
        result.setPromptContent(extraction);
        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 boolean isJsonArray(String json) {
        try {
            new cn.hutool.json.JSONArray(json);
            return true;
        } catch (JSONException e) {
            return false;
        }
    }
}
