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

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.common.constant.MessageRoleConstant;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dbchain.DBChainResult;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dialogue.ImageUrl;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dialogue.Message;
import cn.com.poc.thirdparty.resource.demand.ai.entity.dialogue.MultiContent;
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.JSONObject;
import org.springframework.stereotype.Component;

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

/**
 * @author alex.yao
 * @date 2025/1/14
 */
@Component
public class ImageOCRFunction extends AbstractLargeModelFunction {

    private final String DESC = "该方法是通过OCR获取图片内容，并且根据用户提问对图片内容进行提取总结与分析。";

    private final FunctionLLMConfig functionLLMConfig = new FunctionLLMConfig.FunctionLLMConfigBuilder()
            .name("image_ocr")
            .description(DESC)
            .parameters(new Parameters("object")
                    .addProperties("query", new Properties("string", "用户提问问题"))
                    .addProperties("image_url", new Properties("string", "图片地址")))
            .build();

    private String prompt = "根据用户上传的图片和提出的问题，对图片内容进行提取后根据提出的问题进行总结回答，若是不理解用户提出的问题，则只需要输出图片中的文字内容即可无需对其进行总结等行为，请严格遵守以上内容。";

    private String model = "Doubao-vision-lite-32k";

    @Resource
    private LLMService llmService;

    @Override
    public AbstractFunctionResult<String> doFunction(String content, String identifier, List<DBChainResult> dbChainResults, List<KnowledgeContentResult> knowledgeContentResults) {
        AbstractFunctionResult<String> result = new AbstractFunctionResult<>();
        JSONObject jsonObject = JSONObject.parseObject(content);
        if (!jsonObject.containsKey("query") || !jsonObject.containsKey("image_url")) {
            result.setPromptContent(content);
            result.setFunctionResult(content);
            return result;
        }


        Message systemMessage = new Message();
        systemMessage.setRole(MessageRoleConstant.system);
        systemMessage.setContent(prompt);

        Message userMessage = new Message();

        List<MultiContent> multiContents = new ArrayList<>();
        MultiContent textmultiContent = new MultiContent();
        textmultiContent.setType("text");
        textmultiContent.setText(jsonObject.getString("query"));
        multiContents.add(textmultiContent);
        MultiContent imageMultiContent = new MultiContent();
        imageMultiContent.setType("image_url");
        ImageUrl imageUrl = new ImageUrl();
        imageUrl.setUrl(jsonObject.getString("image_url"));
        imageMultiContent.setImageUrl(imageUrl);
        multiContents.add(imageMultiContent);

        userMessage.setRole(MessageRoleConstant.user);
        userMessage.setContent(multiContents);
        Message[] messages = new Message[]{systemMessage, userMessage};

        LargeModelResponse response = new LargeModelResponse();
        response.setModel(model);
        response.setMessages(messages);
        response.setStream(false);
        response.setUser("Image_OCR");
        LargeModelDemandResult largeModelDemandResult = llmService.chat(response);

        result.setPromptContent(largeModelDemandResult.getMessage());
        result.setFunctionResult(largeModelDemandResult.getMessage());
        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();
    }
}
