yhlxj/netflix-service/src/main/java/com/cyksj/service/midjourney/impl/MidjourneyServiceImpl.java

package com.cyksj.service.midjourney.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.cyksj.common.EnvCommonService;
import com.cyksj.common.exception.BusinessRuntimeException;
import com.cyksj.common.task.GlobalThreadPoolTaskExecutor;
import com.cyksj.common.util.Codec;
import com.cyksj.common.util.J11HttpC;
import com.cyksj.common.util.Jsons;
import com.cyksj.common.util.StringUtil;
import com.cyksj.mapper.MidjourneyUserConversationMapper;
import com.cyksj.mapper.MidjourneyUserMapper;
import com.cyksj.model.dto.BlendDimensions;
import com.cyksj.model.dto.MessageButton;
import com.cyksj.model.entity.MidjourneyAccount;
import com.cyksj.model.entity.MidjourneyUser;
import com.cyksj.model.entity.MidjourneyUserConversation;
import com.cyksj.model.response.SubmitResult;
import com.cyksj.redis.RedisService;
import com.cyksj.service.midjourney.MidjourneyAccountService;
import com.cyksj.service.midjourney.MidjourneyService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.imageio.ImageIO;
import javax.imageio.stream.FileImageOutputStream;
import java.awt.image.BufferedImage;
import java.io.*;
import java.net.URI;
import java.net.URLEncoder;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.security.MessageDigest;
import java.time.ZoneOffset;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

/**
 * @author zwhui
 * @date 2024/4/23 15:44
 */
@Service
@Slf4j
@RequiredArgsConstructor
public class MidjourneyServiceImpl implements MidjourneyService {
    private static final String RELAX_HOST = "http://43.154.230.104:8080";
    private static final String FAST_HOST = "https://aigc.api4midjourney.com/api";
    private static final String FAST_TOKEN = "14a6b469-7cb2-4b99-9238-676c7c5cffef";

    private final MidjourneyUserConversationMapper conversationMapper;

    private static final List<String> progress = List.of("100%");

    private static final List<String> status = List.of("MODAL","CANCEL","FAILURE");

    private final RedisService redisService;

    private final MidjourneyUserMapper midjourneyUserMapper;

    @Value("${midjourney.url}")
    private String midjourneyHost;

    private final EnvCommonService envCommonService;

    private static final GlobalThreadPoolTaskExecutor TASK_EXECUTOR = GlobalThreadPoolTaskExecutor.getInstance();

    private final MidjourneyAccountService midjourneyAccountService;


    @Override
    public MidjourneyUserConversation submitImagine(MidjourneyUser user, String prompt,String botType, List<String> base64Array) throws Exception {
        Map<String, Object> imagineParam = MapUtil.builder(new HashMap<String,Object>())
                .put("prompt", prompt)
                .put("state", user.getId())
                .build();
        if (CollectionUtil.isNotEmpty(base64Array)) {
            imagineParam.put("base64Array", base64Array);
        }
        SubmitResult result = submit(user.getMode(),"imagine", null,imagineParam);
        if (user.getMode() == 2) {
            redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
        }
        return saveConversation(user.getId(), user.getMode(),result,"IMAGINE", StringUtils.EMPTY, botType);
    }

    public MidjourneyUserConversation saveConversation(Long userId,Integer mode,SubmitResult result, String action,String prompt, String botType) throws Exception {
        Long taskId = Long.parseLong(result.getResult());
        Map<String, Object> properties = result.getProperties();
        MidjourneyUserConversation conversation = new MidjourneyUserConversation()
                .setUserId(userId).setTaskId(taskId).setAction(action).setPrompt(prompt)
                .setMode(mode).setStartTime(System.currentTimeMillis()).setProgress("0%")
                .setStatus(result.getCode() == 21 ? "NOT_START" : "IN_PROGRESS").setTaskId(taskId);
        if (MapUtil.isNotEmpty(properties)) {
            conversation.setProperties(Jsons.toJson(properties));
            conversation.setChannelId(properties.get("discordChannelId") == null ? null : Long.parseLong(properties.get("discordChannelId").toString()))
                    .setInstanceId(properties.get("discordInstanceId") == null ? null : Long.parseLong(properties.get("discordInstanceId").toString()));
        }
        if (conversation.getInstanceId() == null) {
            conversation.setInstanceId(result.getInstanceId());
        }
        conversationMapper.insert(conversation);
        return conversation;
    }


    @Override
    public MidjourneyUserConversation submitDescribe(MidjourneyUser user, String botType, String base64) throws Exception {
        Map<String, Object> param = MapUtil.builder(new HashMap<String,Object>())
                .put("state", user.getId())
                .put("base64", base64)
                .build();
        SubmitResult result = submit(user.getMode(),"describe", null, param);
        if (user.getMode() == 2) {
            redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
        }
        return saveConversation(user.getId(), user.getMode(),result,"DESCRIBE", StringUtils.EMPTY, botType);
    }

    @Override
    public MidjourneyUserConversation submitBlend(MidjourneyUser user, BlendDimensions dimensions, String botType, List<String> base64Array) throws Exception {
        Map<String, Object> param = MapUtil.builder(new HashMap<String,Object>())
                .put("base64Array", base64Array)
                .put("state", user.getId())
                .build();
        if (dimensions != null) {
            param.put("dimensions", dimensions);
        }
        SubmitResult result = submit(user.getMode(),"blend", null, param);
        if (user.getMode() == 2) {
            redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
        }
        return saveConversation(user.getId(), user.getMode(),result,"BLEND", StringUtils.EMPTY, botType);
    }

    @Override
    public MidjourneyUserConversation submitModal(MidjourneyUser user, Long taskId, String prompt, String maskBase64) throws Exception {
        Map<String, Object> param = MapUtil.builder(new HashMap<String,Object>())
                .put("taskId", taskId)
                .put("state", user.getId())
                .build();
        if (StringUtils.isNotBlank(prompt)) {
            param.put("prompt", prompt);
        }
        if (StringUtils.isNotBlank(maskBase64)) {
            param.put("maskBase64", maskBase64);
        }
        SubmitResult result = submit(user.getMode(),"modal", null, param);
        if (user.getMode() == 2) {
            redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
        }
        return saveConversation(user.getId(),  user.getMode(),result,"MODAL", StringUtils.EMPTY, StringUtils.EMPTY);
    }

    @Override
    public MidjourneyUserConversation submitShorten(MidjourneyUser user, String botType, String prompt) throws Exception {
        Map<String, Object> param = MapUtil.builder(new HashMap<String,Object>())
                .put("prompt", prompt)
                .put("state", user.getId())
                .build();
        SubmitResult result = submit(user.getMode(),"shorten", null, param);
        if (user.getMode() == 2) {
            redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
        }
        return saveConversation(user.getId(), user.getMode(), result,"SHORTEN", StringUtils.EMPTY, botType);
    }

    /**
     * 恢复次数
     */
    public Long recoverUserLimit(Long id,Integer mode,Long num){
        if (mode == 1){
            num =  redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + id, 1L);
        }
        if (mode == 2){
            if (num != null) {
                num = redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + id, 1L);
            }
        }
        return num;
    }
//    /**
//     * 同步数据库
//     */
//    public void syncUser(Long id,Integer mode,Long num){
//        log.info("同步次数 id:{},mode:{},num:{}",id,mode,num);
//        if (num == null){
//            return;
//        }
//        LambdaUpdateWrapper<MidjourneyUser> wrapper = Wrappers.lambdaUpdate(MidjourneyUser.class)
//                .eq(MidjourneyUser::getId, id);
//        if (mode == 1){
//            wrapper.set(MidjourneyUser::getMjFastNum, num);
//        }
//        if (mode == 2){
//            wrapper.set(MidjourneyUser::getMjRelaxNum, num);
//        }
//        midjourneyUserMapper.update(null, wrapper);
//    }

    @Override
    public SubmitResult submitAction(MidjourneyUser user, Long taskId, String customId, Long num, String botType) throws Exception {
        MidjourneyUserConversation conversation = conversationMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUserConversation.class).eq(MidjourneyUserConversation::getTaskId, taskId).last("limit 1"));
        if (conversation == null) {
            throw BusinessRuntimeException.getInstance("关联任务不存在或已失效");
        }
        if (!user.getMode().equals(conversation.getMode())) {
            throw BusinessRuntimeException.getInstance("当前出图模式与关联任务出图模式不符");
        }
        if (StringUtils.isNotBlank(conversation.getButtons())){
            List<MessageButton> messageButtons = Jsons.parseList(conversation.getButtons(), MessageButton.class);
            messageButtons.forEach(button -> {
                if (button.getCustomId().equals(customId)) {
                    log.info("action customId:{}", customId);
                    button.setStyle(3);
                }
            });
            conversation.setButtons(Jsons.toJson(messageButtons));
        }
        if (customId.contains("BOOKMARK")) {
            conversation.setBookmark(true);
            conversationMapper.updateById(conversation);
            return null;
        }
        Map<String, Object> param = MapUtil.builder(new HashMap<String,Object>())
                .put("taskId", taskId)
                .put("state", user.getId())
                .put("customId", customId)
                .build();
        SubmitResult result = submit(user.getMode(),"action", conversation.getInstanceId(),param);
        if (result.getCode() == 21) {
            // 以上操作有弹窗确认，恢复次数
            recoverUserLimit(user.getId(), user.getMode(),num);
        }else {
            //syncUser(user.getId(), user.getMode(),num);
            if (!customId.contains("upsample") && user.getMode() == 2) {
                redisService.hincr(RedisService.key.MIDJOURNEY_ACCOUNT.getName(), result.getInstanceId().toString(), 1.0);
            }
        }
        saveConversation(user.getId(), user.getMode(), result,"ACTION", StringUtils.EMPTY, botType);
        return result;
    }


    public SubmitResult submit(Integer mode,String action,Long instanceId, Map<String, Object> param) throws Exception {
        String url = "";
        String accountWithMinUsage = "";
        if (mode == 1) {
            param.put("mode", "FAST");
            url = FAST_HOST;
        } else if (mode == 2) {
            url = RELAX_HOST;
            //慢速查询在使用次数最少的账号
            accountWithMinUsage = getAccountWithMinUsage(instanceId);
            param.put("accountFilter",MapUtil.builder(new HashMap<String,Object>())
                    .put("instanceId",accountWithMinUsage).build());
        }
        url = url + getActionUrl(action);
        param.put("notifyHook",midjourneyHost +(EnvCommonService.active.equals(envCommonService.getEnv()) ? "/8081":"/8082") + "/api/applets/midjourney/notifyHook");
        String body = HttpRequest.post(url).body(Jsons.toJson(param)).header("Authorization", FAST_TOKEN).execute().body();
        log.info("action body:{}", body);
        SubmitResult submitResult = Jsons.parseObject(body, SubmitResult.class);
        int code = submitResult.getCode();
        if (code != 1 && code != 21 && code != 22) {
            if (code == 3) {
                if(mode == 2 && StringUtils.isNotBlank(accountWithMinUsage)){
                    redisService.hdel(RedisService.key.MIDJOURNEY_ACCOUNT.getName(),accountWithMinUsage);
                    String finalAccountWithMinUsage = accountWithMinUsage;
                    TASK_EXECUTOR.execute(() -> {
                        MidjourneyAccount account = midjourneyAccountService.getOne(Wrappers.lambdaQuery(MidjourneyAccount.class).eq(MidjourneyAccount::getInstanceId, finalAccountWithMinUsage).last("limit 1"));
                        midjourneyAccountService.updateStatus(account.getId());
                    });
                }
                throw BusinessRuntimeException.getInstance("账号不存在");
            }
            if (code == 4) {
                throw BusinessRuntimeException.getInstance(submitResult.getDescription());
            }
            if (code == 24) {
                throw BusinessRuntimeException.getInstance("prompt包含敏感词");
            }
            log.error("action:" + action + " error message:" + submitResult.getDescription());
            throw BusinessRuntimeException.getInstance("队列已满，请稍后尝试");
        }
        // 增加使用次数
        if (StringUtils.isNotBlank(accountWithMinUsage) && mode == 2) {
            submitResult.setInstanceId(Long.valueOf(accountWithMinUsage));
        }
        return submitResult;
    }

    public String getAccountWithMinUsage(Long instanceId) {
        String key = RedisService.key.MIDJOURNEY_ACCOUNT.getName();
        try {
            // 获取所有账号ID和使用次数
            Map<Object, Object> accountsUsage = redisService.hmget(key);

            if (accountsUsage == null || accountsUsage.isEmpty()) {
                accountsUsage = midjourneyAccountService.list(Wrappers.lambdaQuery(MidjourneyAccount.class).eq(MidjourneyAccount::getStatus,Boolean.TRUE)).stream().collect(Collectors.toMap(k -> k.getInstanceId().toString(), v -> 0));
                // 保存所有账号ID和使用次数
                redisService.hmset(key, accountsUsage);
            }

            // 找到使用次数最少的账号ID
            String minAccountId = null;
            if (instanceId != null) {
                minAccountId = instanceId.toString();
            }else {
                int minUsage = Integer.MAX_VALUE;
                for (Map.Entry<Object, Object> entry : accountsUsage.entrySet()) {
                    int usage = Integer.parseInt(entry.getValue().toString());
                    if (usage < minUsage) {
                        minUsage = usage;
                        minAccountId = entry.getKey().toString();
                    }
                }
            }
            return minAccountId;
        } catch (Exception e) {
            throw BusinessRuntimeException.getInstance("获取账号失败");
        }
    }
    private String getActionUrl(String action) {
         switch (action) {
            case "imagine":
                return "/mj/submit/imagine";
            case "action":
                return "/mj/submit/action";
            case "describe":
                return "/mj/submit/describe";
            case "blend":
                return"/mj/submit/blend";
            case "modal":
                return "/mj/submit/modal";
            case "shorten":
                return "/mj/submit/shorten";
            default:  throw new IllegalArgumentException("Unknown action: " + action);
        }
    }
    @Override
    public List<MidjourneyUserConversation> listConversationByIds(Integer mode, List<Long> ids) {
        List<MidjourneyUserConversation> list = new ArrayList<>();
        ids.forEach(id ->{
            String queryKey = RedisService.key.MIDJOURNEY_QUERY.getName();
            Long count = redisService.incr(queryKey + id, 1L);
            if (count%5 == 0) {
                try {
                    listByIds(mode,List.of(id)).forEach(json ->{
                        JSONObject jsons = JSONUtil.parseObj(json);
                        MidjourneyUserConversation conversation = JSONUtil.toBean(jsons, MidjourneyUserConversation.class);
                        if (StringUtil.isNotBlank(conversation.getImageUrl())) {
                            conversation.setImageUrl(conversation.getImageUrl().replace("cdn.discordapp.com", "mj.galaxydvd.com"));
                        }
                        conversation.setUserId(jsons.getLong("state"));
                        conversation.setTaskId(jsons.getLong("id"));
                        list.add(conversation);
                        TASK_EXECUTOR.execute(() -> {
                            try {
                                sync(conversation);
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }
                        });
                    });
                } catch (Exception e) {
                    throw new RuntimeException(e);
                }
            }else {
                list.add((MidjourneyUserConversation) redisService.get(RedisService.key.MIDJOURNEY_CONVERSATION.getName() + id));
            }
        });
        return list;
    }


    public void sync(MidjourneyUserConversation conversation) throws IOException {
        if ((StringUtils.isNotBlank(conversation.getProgress()) && progress.contains(conversation.getProgress())) || status.contains(conversation.getStatus())) {
            log.info("同步任务：{},进度:{}",conversation.getTaskId(),conversation.getProgress());
            Long userId = conversation.getUserId();
            MidjourneyUserConversation dbConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUserConversation.class).eq(MidjourneyUserConversation::getTaskId, conversation.getTaskId())
                    .eq(MidjourneyUserConversation::getUserId, userId).orderByDesc(MidjourneyUserConversation::getId).last("limit 1"));
            if (dbConversation != null) {
                if ("SUCCESS".equals(dbConversation.getStatus()) || "FAILURE".equals(dbConversation.getStatus()) || "MODAL".equals(dbConversation.getStatus())) {
                    return;
                }
                conversation.setId(dbConversation.getId());
                conversationMapper.updateById(conversation);
                TASK_EXECUTOR.execute(() -> {
                    try {
                        String imageUrl = conversation.getImageUrl();
                        if (StringUtil.isNotBlank(imageUrl)) {
                            conversation.setImageUrl(uploadPic(imageUrl, "conversation"+dbConversation.getId()));
                            conversationMapper.updateById(conversation);
                        }
                    } catch (IOException e) {
                        log.error("上传图片失败",e);
                    }
                });
                //失败返还次数
                if ("FAILURE".equals(conversation.getStatus())){
                    Object num = redisService.get(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + userId);
                    LambdaUpdateWrapper<MidjourneyUser> wrapper = Wrappers.lambdaUpdate(MidjourneyUser.class)
                            .eq(MidjourneyUser::getId, userId);
                    Integer mode = dbConversation.getMode();
                    if (mode == 1){
                        num =  redisService.incr(RedisService.key.MIDJOURNEY_FAST_LIMIT.getName() + userId, 1L);
                        wrapper.set(MidjourneyUser::getMjFastNum, num);
                    }
                    if (mode == 2){
                        if (num != null) {
                            num = redisService.incr(RedisService.key.MIDJOURNEY_RELAX_LIMIT.getName() + userId, 1L);
                            wrapper.set(MidjourneyUser::getMjRelaxNum, num);
                        }
                    }
                    midjourneyUserMapper.update(null, wrapper);
                }
            }
        }
    }


    private static String uploadPic(String url, String prefix) throws IOException {
        byte[] body = HttpUtil.downloadBytes(url);
        Path tempFile = Files.createTempFile(prefix, ".png");
        File file = tempFile.toFile();
        try (FileImageOutputStream imageOutput = new FileImageOutputStream(file)) {
            imageOutput.write(body, 0, body.length);
        }
        Map<String, Object> paramMap = new HashMap<>();
        paramMap.put("file", file);
        String json = HttpUtil.post("https://files.liuliangbang.vip/pic/ups", paramMap);
        log.info("上传图片结果：url:{},json:{}",url, json);
        FileUtil.del(file);
        return new JSONObject(json).getJSONObject("value").getJSONArray("saved").getJSONObject(0).getJSONObject("info").getStr("cdnUrl");
    }

    private static List<String> uploadBase64Pic(List<String> base64Array) throws IOException {
        List<String> list = new ArrayList<>();
        for (String base64 : base64Array) {
            String fileExt = "jpeg";
            if (base64.contains(";")){
                fileExt = base64.split(";")[0].split("/")[1];
                base64 = base64.split(",")[1];
            }
            byte[] body = Base64.getDecoder().decode(base64);
            Path tempFile = Files.createTempFile(UUID.randomUUID().toString(), "." + fileExt);
            try (FileImageOutputStream imageOutput = new FileImageOutputStream(tempFile.toFile())) {
                imageOutput.write(body, 0, body.length);
            }
            Map<String, Object> paramMap = new HashMap<>();
            paramMap.put("file", tempFile.toFile());
            String json = HttpUtil.post("https://files.liuliangbang.vip/pic/ups", paramMap);
            list.add(new JSONObject(json).getJSONObject("value").getJSONArray("saved").getJSONObject(0).getJSONObject("info").getStr("cdnUrl"));
        }
        return list;
    }
    public JSONArray listByIds(Integer mode,List<Long> ids) throws Exception {
        Map<String, Object> param = MapUtil.builder(new HashMap<String, Object>()).put("ids", ids).build();
        String body = HttpRequest.post((mode == 1 ? FAST_HOST : RELAX_HOST) + "/mj/task/list-by-condition").header("Authorization", FAST_TOKEN).body(Jsons.toJson(param)).execute().body();
        log.info("listByIds body:{}", body);
        return JSONUtil.parseArray(body);
    }

    @Override
    public MidjourneyUserConversation cancelConversation(MidjourneyUser user, Long id) {
        MidjourneyUserConversation conversation = conversationMapper.selectOne(Wrappers.lambdaQuery(MidjourneyUserConversation.class).eq(MidjourneyUserConversation::getId, id));
        if (conversation == null) {
            throw BusinessRuntimeException.getInstance("会话不存在");
        }
        if (!Objects.equals(conversation.getUserId(), user.getId())) {
            throw BusinessRuntimeException.getInstance("不是您的会话");
        }
        conversation.setStatus("CANCEL");
        conversationMapper.updateById(conversation);
        TASK_EXECUTOR.execute(() -> {
            cancel(user.getMode(),conversation.getTaskId());
        });
        return conversation;
    }

    public void cancel(Integer mode,Long taskId) {
        String body = HttpUtil.post(mode == 1 ? FAST_HOST : RELAX_HOST + "/mj/task/"+taskId+"/cancel",MapUtil.empty());
        log.info("cancel body:{}", body);
        JSONObject jsonObject = JSONUtil.parseObj(body);
        if (jsonObject.getInt("code") != 1) {
            throw BusinessRuntimeException.getInstance("cancel error message:" + jsonObject.getStr("description"));
        }
    }

    @Override
    public void notifyHook(String json) {
        log.info("notifyHook:{}", json);
        JSONObject jsons = JSONUtil.parseObj(json);
        MidjourneyUserConversation conversation = JSONUtil.toBean(jsons, MidjourneyUserConversation.class);
        conversation.setUserId(jsons.getLong("state"));
        conversation.setTaskId(jsons.getLong("id"));
        if (StringUtil.isNotBlank(conversation.getImageUrl())) {
            conversation.setImageUrl(conversation.getImageUrl().replace("cdn.discordapp.com", "mj.galaxydvd.com"));
        }
        redisService.set(RedisService.key.MIDJOURNEY_CONVERSATION.getName() + conversation.getTaskId(), conversation,RedisService.key.MIDJOURNEY_CONVERSATION.getTimeout());
        TASK_EXECUTOR.execute(() -> {
            try {
                sync(conversation);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        });
    }
}