package com.sonic.cow.client.impl;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.byteplus.ark.runtime.model.images.generation.GenerateImagesRequest;
import com.byteplus.ark.runtime.model.images.generation.ImagesResponse;
import com.byteplus.ark.runtime.model.images.generation.ResponseFormat;
import com.byteplus.ark.runtime.service.ArkService;
import com.google.common.collect.Maps;
import com.mashape.unirest.http.HttpResponse;
import com.mashape.unirest.http.JsonNode;
import com.mashape.unirest.http.Unirest;
import com.mashape.unirest.http.exceptions.UnirestException;
import com.sonic.cow.client.Seedream4GenImageClient;
import com.sonic.cow.client.TextGenImageClient;
import com.sonic.cow.utils.LimitUtils;
import com.sonic.cow.utils.MD5Util;
import com.sonic.cow.utils.RedisKeyUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.Map;

@Slf4j
@Service
public class Seedream4ImageClientImpl implements Seedream4GenImageClient {

    @Value("${volcengine.seedream4Model}")
    private String seedream4Model;
    @Value("${volcengine.apiKey}")
    private String apiKey;


    @Autowired
    private ArkService volcengineArkService;
    @Autowired
    private LimitUtils limitUtils;
    @Autowired
    private RedisKeyUtils redisKeyUtils;

    @Override
    public String genImage(String imageUrl, String prompt) throws Exception {
        try {
            // 构建请求体 JSON
            Map<Object, Object> requestBody = Maps.newHashMap();
            requestBody.put("model", seedream4Model);
            requestBody.put("prompt", prompt);
            if (StringUtils.isNotBlank(imageUrl)) {
                requestBody.put("image", imageUrl);
            }
            requestBody.put("sequential_image_generation", "disabled");
            requestBody.put("response_format", "b64_json");
            requestBody.put("size", "1440x2560");
            requestBody.put("stream", false);
            requestBody.put("watermark", false);

            // 发送 POST 请求
            HttpResponse<JsonNode> response = Unirest.post("https://ark.ap-southeast.bytepluses.com/api/v3/images/generations")
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + apiKey) // 拼接 Bearer 令牌
                    .body(JSON.toJSONString(requestBody)) // 设置请求体
                    .asJson(); // 以 JSON 格式接收响应
            // 处理响应
            if (response.getStatus() == 200) {
                log.info("请求成功，响应结果：{}", response.getBody().toString());
                // 可从 response.getBody() 中解析生成的图片 URL 等信息
                String base64 = response.getBody().getObject().getJSONArray("data").getJSONObject(0).getString("b64_json");
                return base64;
            } else {
                log.info("请求失败，状态码：：{}", response.getStatus());
                log.info("错误信息：：{}", response.getBody());
                throw new Exception(JSON.toJSONString(response.getBody()));
            }
        } catch (Exception e) {
           log.error("===> seedream4 genImage error: ", e);
            String md5Prompt =  MD5Util.digest(prompt);
            String imageGenImageErrorEnumKey = redisKeyUtils.imageGenImageErrorEnumKey(md5Prompt);
            int errorNum = limitUtils.defaultLimitCheckReturnCount(imageGenImageErrorEnumKey, 10 * 60);
            log.info("seedream4 imageGenImage errorNum:{}", errorNum);
            //最多重试2次
            if (errorNum <= 2) {
                log.info("seedream4 imageGenImage 重试次数:{}", errorNum);
                //调用图生图  角色特征保持方法
                return genImage(imageUrl, prompt);
            }
            throw new Exception(e.getMessage());
        } finally {
            volcengineArkService.shutdownExecutor();
        }
    }

    public static void main(String[] args) {
        // 替换为你的 ARK API 密钥
        String arkApiKey = "4075846b-cb9e-45f6-bf59-ab0a94745d9b";

        try {
            // 构建请求体 JSON
//            String requestBody = "{" +
//                    "\"model\": \"ep-20250911184318-89dxw\"," +
//                    "\"prompt\": \"海滩，穿着比基尼，悠闲散步，手里拿着椰子，张开双手，享受海风，开怀大笑\"," +
//                    "\"image\": \"https://hhb.crushlevel.ai/dev/role/17562030935666537.jpg\"," +
//                    "\"sequential_image_generation\": \"disabled\"," +
////                    "\"response_format\": \"b64_json\"," +
//                    "\"response_format\": \"url\"," +
//                    "\"size\": \"2K\"," +
//                    "\"stream\": false," +
//                    "\"watermark\": true" +
//                    "}";

            Map<Object, Object> requestBody = Maps.newHashMap();
            requestBody.put("model", "ep-20250911184318-89dxw");
            requestBody.put("prompt", "海滩，穿着比基尼，悠闲散步，手里拿着椰子，张开双手，享受海风，开怀大笑");
//            requestBody.put("image", "https://hhb.crushlevel.ai/dev/role/17562030935666537.jpg");
            requestBody.put("sequential_image_generation", "disabled");
            requestBody.put("response_format", "url");
//            requestBody.put("response_format", "b64_json");
            requestBody.put("size", "1440x2560");
            requestBody.put("stream", false);
            requestBody.put("watermark", false);

            // 发送 POST 请求
            HttpResponse<JsonNode> response = Unirest.post("https://ark.ap-southeast.bytepluses.com/api/v3/images/generations")
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer " + arkApiKey) // 拼接 Bearer 令牌
                    .body(JSON.toJSONString(requestBody)) // 设置请求体
                    .asJson(); // 以 JSON 格式接收响应
            // 处理响应
            if (response.getStatus() == 200) {
                System.out.println("请求成功，响应结果：");
                System.out.println(response.getBody().toString());
                // 可从 response.getBody() 中解析生成的图片 URL 等信息
                System.out.println("图片 URL：" + response.getBody().getObject().getJSONArray("data").getJSONObject(0).getString("url"));
            } else {
                System.out.println("请求失败，状态码：" + response.getStatus());
                System.out.println("错误信息：" + response.getBody());
            }

        } catch (Exception e) {
            e.printStackTrace();
        } finally {

        }
    }
}
