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

import cn.com.poc.agent_application.entity.KnowledgeContentResult;
import cn.com.poc.agent_application.entity.Variable;
import cn.com.poc.common.service.RedisService;
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.com.yict.framemax.core.i18n.I18nMessageException;
import cn.hutool.core.collection.ListUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONException;
import cn.hutool.json.JSONObject;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Component
public class MemoryVariableWriterFunction extends AbstractLargeModelFunction {

    private final String DESC = "该方法参考key字段给出的内容键，来保存对应的内容值，该方法不可进行查询";

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig
            .FunctionLLMConfigBuilder()
            .name("memory_variable_writer")
            .description(DESC)
            .parameters(new Parameters("array")
                    .addProperties("key", new Properties("string", "内容键"))
                    .addProperties("value", new Properties("string", "内容值"))
            ).build();


    @Resource
    private RedisService redisService;

    @Override
    public AbstractFunctionResult<String> doFunction(String content, String identifier, List<DBChainResult> dbChainResults, List<KnowledgeContentResult> knowledgeContentResults) {
        AbstractFunctionResult<String> result = new AbstractFunctionResult<>();
        if (StringUtils.isBlank(content) || StringUtils.isBlank(identifier)) {
            result.setPromptContent("FAIL");
            result.setFunctionResult("FAIL");
            return result;
        }

        String contentKey = MemoryVariableWriterConstants.REDIS_PREFIX + identifier;
        // 创建 JSONObject 对象
        Map<String, Object> map = new HashMap<>();
        if (redisService.hasKey(contentKey)) {
            Map<Object, Object> hmget = redisService.hmget(contentKey);
            for (Map.Entry<Object, Object> entry : hmget.entrySet()) {
                if (entry.getKey() instanceof String) {
                    String tempKey = (String) entry.getKey();
                    map.put(tempKey, entry.getValue());
                }
            }
        }
        if (isJsonArray(content)) {
            JSONArray jsonArray = new JSONArray(content);
            for (int i = 0; i < jsonArray.size(); i++) {
                setMap(jsonArray.getJSONObject(i), map);
            }
        } else {
            setMap(new JSONObject(content), map);
        }
        redisService.hmset(contentKey, map);

        result.setPromptContent("SUCCESS");
        result.setFunctionResult("SUCCESS");
        return result;
    }

    private static void setMap(JSONObject jsonObject, Map<String, Object> result) {
        String key = jsonObject.getStr("key");
        String value = jsonObject.getStr("value");
        result.put(key, value);
    }

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

    @Override
    public List<String> getLLMConfig() {
        throw new I18nMessageException("exception/this.method.is.not.supported");
    }

    @Override
    public List<String> getLLMConfig(List<Variable> variableStructure) {
        List<String> enumList = new ArrayList<>();
        for (Variable variable : variableStructure) {
            enumList.add(variable.getKey());
        }
        this.functionLLMConfig.getFunction().getParameters().addProperties("key", new Properties("string", "内容键", enumList));
        return ListUtil.toList(JsonUtils.serialize(functionLLMConfig));
    }


    private boolean isJsonArray(String json) {
        try {
            new JSONArray(json);
            return true;
        } catch (JSONException e) {
            return false;
        }
    }

}
