package com.sonic.cow.client.impl;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import com.sonic.cow.client.ImageGenImageClient;
import com.sonic.cow.utils.LimitUtils;
import com.sonic.cow.utils.MD5Util;
import com.sonic.cow.utils.RedisKeyUtils;
import com.volcengine.service.visual.IVisualService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 文档地址：https://www.volcengine.com/docs/85128/1584306#YLSQPqOs
 * SDK说明：https://www.volcengine.com/docs/6444/1340578
 * 实例：https://github.com/volcengine/volc-sdk-java/blob/main/example/src/main/java/com/volcengine/example/visual/CVProcessDemo.java
 */
@Slf4j
@Service
public class ImageToImageClientImpl implements ImageGenImageClient {

    @Autowired
    private IVisualService visualService;
    @Autowired
    private LimitUtils limitUtils;
    @Autowired
    private RedisKeyUtils redisKeyUtils;

    /**
     * 可重试的code
     */
    private static final List<String> RETRIED_CODE = Lists.newArrayList("50511", "50429", "50430", "50500", "50501");


    @Override
    public String imageGenImage(String imageUrl, String prompt) throws Exception {
        try {
            Long startTime = System.currentTimeMillis();

            JSONObject req = new JSONObject();
            //下面三个必填参数
            req.put("req_key", "i2i_portrait_photo");
            req.put("image_input", imageUrl);
            req.put("prompt", prompt);

            //下面参数非必填参数，可自行选择是否添加
            req.put("gpen", "0.4");
            req.put("skin", "0.3");
            req.put("skin_unifi", 0);
            //图片比例
            req.put("width", "720");
            req.put("height", "1280");
            req.put("gen_mode", "creative");
            req.put("seed", "-1");
            req.put("return_url", false);

            Object response = visualService.cvProcess(req);
            JSONObject jsonObject = JSONObject.parseObject(JSON.toJSONString(response));
            log.info("imageToImage 图生成图片结果：{}", jsonObject);
            Long endTime = System.currentTimeMillis();
            log.info("imageToImage 图生成图片所需要的时间：{}", (endTime - startTime) / 1000);
            try {
                //睡眠个3秒，等待资源释放，下一个请求才不会报并发限制
                Thread.sleep(3000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            return jsonObject.getJSONObject("data").getJSONArray("binary_data_base64").get(0).toString();
        } catch (Exception e) {
            log.error("imageGenImage error:", e);
//            String message = e.getMessage();
//            JSONObject errorObject = JSONObject.parseObject(message);
//            String code = errorObject.getString("code");
            String md5Prompt =  MD5Util.digest(prompt);
            String imageGenImageErrorEnumKey = redisKeyUtils.imageGenImageErrorEnumKey(md5Prompt);
            int errorNum = limitUtils.defaultLimitCheckReturnCount(imageGenImageErrorEnumKey, 10 * 60);
            log.info("imageGenImage errorNum:{}", errorNum);
            //最多重试5次
            if (errorNum <= 5) {
                log.info("imageGenImage 重试次数:{}", errorNum);
                //调用图生图  角色特征保持方法
                return imageGenImageV2(imageUrl, prompt);
            }
            throw new Exception(e.getMessage());
        }
    }

    @Override
    public String imageGenImageV2(String imageUrl, String prompt) throws Exception {
        try {
            Long startTime = System.currentTimeMillis();

            JSONObject req = new JSONObject();
            //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
            req.put("req_key", "seed3l_single_ip");
            req.put("image_urls", new String[]{imageUrl});
            req.put("prompt", prompt);
            req.put("seed", "-1");
            req.put("width", "720");
            req.put("height", "1280");
            req.put("gen_mode", "creative");
            req.put("return_url", false);

            Object response = visualService.cvProcess(req);
            JSONObject jsonObject = JSONObject.parseObject(JSON.toJSONString(response));
            log.info("imageToImage 图生成图片结果：{}", jsonObject);
            Long endTime = System.currentTimeMillis();
            log.info("imageGenImageV2 图生成图片所需要的时间：{}", (endTime - startTime) / 1000);
            try {
                //睡眠个3秒，等待资源释放，下一个请求才不会报并发限制
                Thread.sleep(3000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            return jsonObject.getJSONObject("data").getJSONArray("binary_data_base64").get(0).toString();
        } catch (Exception e) {
            log.error("imageGenImageV2 error:", e);
//            String message = e.getMessage();
//            JSONObject errorObject = JSONObject.parseObject(message);
//            String code = errorObject.getString("code");
            String md5Prompt =  MD5Util.digest(prompt);
            String imageGenImageErrorEnumKey = redisKeyUtils.imageGenImageErrorEnumKey(md5Prompt);
            int errorNum = limitUtils.defaultLimitCheckReturnCount(imageGenImageErrorEnumKey, 10 * 60);
            log.info("imageGenImageV2 errorNum:{}", errorNum);
            //最多重试5次
            if (errorNum <= 5) {
                log.info("imageGenImageV2 重试次数:{}", errorNum);
                //调用图生图  人物写真方法
                return imageGenImage(imageUrl, prompt);
            }
            throw new Exception(e.getMessage());
        }
    }

    @Override
    public String imageGenImageSubmitTask(String imageUrl, String prompt) {
        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "i2i_portrait_photo");
        req.put("image_input", imageUrl);
        req.put("prompt", prompt);
        req.put("width", "720");
        req.put("height", "1280");
        try {
            Long startTime = System.currentTimeMillis();
            Object response = visualService.cvSync2AsyncSubmitTask(req);
            JSONObject jsonObject = JSONObject.parseObject(JSON.toJSONString(response));
            Long endTime = System.currentTimeMillis();
            log.info("imageToImage 图生成图片提交任务结果所需要的时间：{}", (endTime - startTime) / 1000);
            return jsonObject.getJSONObject("data").getString("task_id");
        } catch (Exception e) {
            log.error("imageGenImageSubmitTask error:", e);
            return "";
        }
    }

    @Override
    public String imageGenImageGetTaskResult(String taskId) {
        JSONObject req = new JSONObject();
        //请求Body(查看接口文档请求参数-请求示例，将请求参数内容复制到此)
        req.put("req_key", "i2i_portrait_photo");
        req.put("task_id", "14366847939889876549");
//        req.put("req_json","");
        try {
            Long startTime = System.currentTimeMillis();

            Object response = visualService.cvSync2AsyncGetResult(req);
            JSONObject jsonObject = JSONObject.parseObject(JSON.toJSONString(response));

            Long endTime = System.currentTimeMillis();
            log.info("imageToImage 图生成图片获取任务结果所需要的时间：{}", (endTime - startTime) / 1000);
            return jsonObject.getJSONObject("data").getJSONArray("binary_data_base64").get(0).toString();
        } catch (Exception e) {
            log.error("imageGenImageSubmitTask error:", e);
            return "";
        }
    }
}
