package com.heyu.api.alibaba;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.GenerationUsage;
import com.alibaba.dashscope.aigc.multimodalconversation.*;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.MultiModalMessage;
import com.alibaba.dashscope.common.Role;
import com.alibaba.fastjson.JSON;
import com.heyu.api.alibaba.bailian.AlibabaOCREnums;
import com.heyu.api.alibaba.resp.ModelResult;
import com.heyu.api.data.utils.StringUtils;
import lombok.extern.slf4j.Slf4j;

import java.util.*;


@Slf4j
public class LLMUtils {

    public final static String apiKey = "sk-ef6213245c3648ea81f2e4a8ccd34d75";

    public final static String
            prompt = "# 角色\n"
            + "你是一个语言翻译专家，能将用户输入的内容进行翻译\n" +
            "# 任务说明\n" +
            "翻译成 ";
    ;


    public static void main(String[] args) {

        //ModelResult modelResult = callBaiLian("开通阿里云百炼：使用阿里云主账号前往阿里云百炼控制台，如果页面顶部显示以下消息，您需要开通阿里云百炼的模型服务，以获得免费额度。如果未显示该消息，则表示您已经开通。",prompt);

        //System.out.println(JSON.toJSON(modelResult));

        callOcr("https://heyuoss.oss-cn-shanghai.aliyuncs.com/test/ccc.jpg","",AlibabaOCREnums.text_recognition );

    }

    public static ModelResult callBaiLian(String content, String prompt) {
        ModelResult modelResult = new ModelResult();
        try {
            Date startDate = new Date();
            Generation gen = new Generation();
            Message systemMsg = Message.builder()
                    .role(Role.SYSTEM.getValue())
                    .content(prompt)
                    .build();
            Message userMsg = Message.builder()
                    .role(Role.USER.getValue())
                    .content(content)
                    .build();
            GenerationParam param = GenerationParam.builder()
                    // 若没有配置环境变量，请用阿里云百炼API Key将下行替换为：.apiKey("sk-xxx")
                    .apiKey("sk-ef6213245c3648ea81f2e4a8ccd34d75")
                    // 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                    .model("qwen-plus")
                    .messages(Arrays.asList(systemMsg, userMsg))
                    .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                    .build();

            GenerationResult generationResult = gen.call(param);
            String resp = generationResult.getOutput()
                    .getChoices().get(0).
                    getMessage().getContent();


            modelResult.setResult(resp);

            Date endDate = new Date();

            GenerationUsage generationUsage = generationResult.getUsage();

            modelResult.setTokens(generationUsage.getTotalTokens());
            modelResult.setStartTime(startDate);
            modelResult.setEndTime(endDate);

            modelResult.setExet(endDate.getTime() - startDate.getTime());
            return modelResult;
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            log.info("callBaiLian content :{}, callBaiLian modelResult:{},prompt:{}", content, JSON.toJSONString(modelResult), prompt);
        }

        return null;
    }




    /**
     * https://heyuoss.oss-cn-shanghai.aliyuncs.com/test/ccc.jpg
     */
    public static ModelResult callOcr(String image,String imageBase64, AlibabaOCREnums alibabaOCREnums) {
        ModelResult modelResult = new ModelResult();
        try {
            Date startDate = new Date();
            MultiModalConversation conv = new MultiModalConversation();
            Map<String, Object> map = new HashMap<>();
            if(StringUtils.isNotEmpty(image)){
                map.put("image", image);
            }else{
                map.put("image", "data:image/jpeg;base64," + imageBase64);
            }
            // 输入图像的最大像素阈值，超过该值图像会按原比例缩小，直到总像素低于max_pixels
            map.put("max_pixels", "2073600");//限制图片大小 1920x1080
            // 输入图像的最小像素阈值，小于该值图像会按原比例放大，直到总像素大于min_pixels
            map.put("min_pixels", "3136");
            // 开启图像自动转正功能
            map.put("enable_rotate", true);
            // 配置内置的OCR任务
            OcrOptions ocrOptions = OcrOptions.builder()
                    .task(OcrOptions.Task.TEXT_RECOGNITION)
                    .build();
            MultiModalMessage userMessage = MultiModalMessage.builder().role(Role.USER.getValue())
                    .content(Arrays.asList(
                            map,
                            // 当ocr_options中的task字段设置为表格解析时，模型会以下面text字段中的内容作为Prompt，不支持用户自定义
                            Collections.singletonMap("text", alibabaOCREnums.getPrompt()))).build();
            MultiModalConversationParam param = MultiModalConversationParam.builder()
                    // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                    .apiKey(apiKey)
                    .model(alibabaOCREnums.getModel())
                    .message(userMessage)
                    .ocrOptions(ocrOptions)
                    .build();

            MultiModalConversationResult result = conv.call(param);
            String content = result.getOutput().getChoices().get(0).getMessage().getContent().get(0).get("text") + "";

            Date endDate = new Date();
            MultiModalConversationUsage generationUsage = result.getUsage();

            modelResult.setTokens(generationUsage.getTotalTokens());
            modelResult.setStartTime(startDate);
            modelResult.setEndTime(endDate);
            modelResult.setResult(content);
            modelResult.setExet(endDate.getTime() - startDate.getTime());
            return modelResult;
        } catch (Exception e) {
            log.error("callOcr image error :{}", e.getMessage());
        } finally {
            log.info(" callOcr modelResult :{}", JSON.toJSONString(modelResult));
        }
        return null;
    }


}