yhlxj/luma/src/main/java/com/yhlxj/service/midjourney/impl/LumaServiceImpl.java

package com.yhlxj.service.midjourney.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
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.exception.BusinessRuntimeException;
import com.cyksj.common.task.GlobalThreadPoolTaskExecutor;
import com.cyksj.common.util.Jsons;
import com.cyksj.common.util.StringUtil;
import com.cyksj.model.dto.SubmitVideoDTO;
import com.yhlxj.dao.mapper.GroupsRelationMapper;
import com.yhlxj.dao.mapper.midjourney.LumaUserConversationMapper;
import com.yhlxj.dao.mapper.midjourney.LumaUserMapper;
import com.yhlxj.dao.model.entity.GroupsRelation;
import com.yhlxj.dao.model.entity.LumaUser;
import com.yhlxj.dao.model.entity.LumaUserConversation;
import com.yhlxj.dao.model.response.SubmitResult;
import com.yhlxj.dao.model.response.VideoTask;
import com.yhlxj.redis.RedisService;
import com.yhlxj.service.common.EnvCommonService;
import com.yhlxj.service.midjourney.LumaService;
import com.yhlxj.util.COSUtil;
import com.yhlxj.util.VideoCoverExtractor;
import com.yhlxj.web.wss.WssSession;
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 java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author zwhui
 * @date 2024/4/23 15:44
 */
@Service
@Slf4j
@RequiredArgsConstructor
public class LumaServiceImpl implements LumaService {
    private static final String LUMA_BASE_URL = "https://api.bltcy.ai";
    private static final String AUTH = "Bearer sk-YTg6N3QH6PLTFLbM9291B1Ae12754bC4B26225F1B95f442c";

    private final RedisService redisService;

    private final LumaUserMapper lumaUserMapper;

    @Value("${luma.url}")
    private String lumaHost;

    private final EnvCommonService envCommonService;

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

    private final LumaUserConversationMapper conversationMapper;

    private final GroupsRelationMapper groupsRelationMapper;

    private static final String VIDEO_NOTIFY_URL = "/api/applets/luma/video/notify";

    //webSocket
    public static Map<Long, WssSession> WSS_SESSION_MAP = new ConcurrentHashMap<>();

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

    private static final List<String> status = List.of("pending", "processing");

    private static final List<String> FINISH_STATUS = List.of("completed", "failed");

    @Override
    public LumaUser getUser(String userToken) {
        if(StringUtils.isBlank(userToken)){
            throw BusinessRuntimeException.getInstance("10001","请重新至官网登录");
        }
        LumaUser lumaUser = lumaUserMapper.selectOne(Wrappers.lambdaQuery(LumaUser.class)
                .eq(LumaUser::getUserToken, userToken).last(" limit 1"));
        if (lumaUser == null){
            throw BusinessRuntimeException.getInstance("10001","会话已过期,请重新登录");
        }
        if(!redisService.hasKey(RedisService.key.LUMA_NUM_LIMIT.getName() + lumaUser.getId())){
            redisService.set(RedisService.key.LUMA_NUM_LIMIT.getName() + lumaUser.getId(), lumaUser.getNum());
        }

        Long relationId = lumaUser.getRelationId();
        GroupsRelation relation = groupsRelationMapper.selectById(relationId);
        lumaUser.setAqType(relation.getAqType());
        return lumaUser;
    }

    public LumaUserConversation saveConversation(Long userId,SubmitResult result,String prompt) throws Exception {
        String taskId = result.getId();
        LumaUserConversation conversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class)
                .eq(LumaUserConversation::getTaskId, result.getId())
                .last("limit 1"));
        if (conversation == null){
            conversation = new LumaUserConversation();
            conversation.setUserId(userId);
            conversation.setTaskId(taskId);
            conversation.setPrompt(prompt);
            conversation.setStartTime(System.currentTimeMillis());
            conversation.setProgress("0%");
            conversation.setStatus("pending");
            conversation.setTaskId(taskId);
        }
        if (conversation.getId() == null) {
            conversationMapper.insert(conversation);
        }else {
            conversationMapper.updateById(conversation);
        }
        return conversation;
    }

    public SubmitResult submit(LumaUser user, Map<String, Object> param, Long num) throws Exception {
        String url = LUMA_BASE_URL + "/luma/generations";
        HttpRequest request = HttpRequest.post(url).body(Jsons.toJson(param));
        request.header("Authorization", AUTH);

        String body = request.setConnectionTimeout(10000).execute().body();
        log.info("return body:{}", body);
        SubmitResult submitResult = Jsons.parseObject(body, SubmitResult.class);
        if (submitResult.getId() == null) {
            throw BusinessRuntimeException.getInstance("请输入正确的生成视频信息");
        }
        return submitResult;
    }

    @Override
    public SubmitResult submitVideo(LumaUser user,SubmitVideoDTO videoDTO, Long num) throws Exception {
        Map<String, Object> videoParam = MapUtil.builder(new HashMap<String, Object>())
                .put("aspect_ratio", "16:9")
                .build();
        String prompt = videoDTO.getPrompt();
        Boolean expandPrompt = videoDTO.getExpandPrompt();
        String picture = videoDTO.getPicture();
        String imageEndUrl = videoDTO.getImageEndUrl();
        if (StrUtil.isNotBlank(prompt)) {
            videoParam.put("user_prompt", prompt);
        }
        if (expandPrompt != null) {
            videoParam.put("expand_prompt", expandPrompt);
        }
        if (StrUtil.isNotBlank(picture)) {
            videoParam.put("image_url", picture);
        }
        if (StrUtil.isNotBlank(imageEndUrl)) {
            videoParam.put("image_end_url", imageEndUrl);
        }
        videoParam.put("notify_hook", lumaHost + (EnvCommonService.active_pre.equals(envCommonService.getEnv()) ? "/8083" : "/8085") + VIDEO_NOTIFY_URL);
        SubmitResult result = submit(user, videoParam, num);
        LumaUserConversation lumaUserConversation = saveConversation(user.getId(), result, prompt);
        result.setRid(lumaUserConversation.getId());
        return result;
    }


    @Override
    public SubmitResult getVideoTaskInfoByTaskId(LumaUser user, Long rid) throws Exception {
        LumaUserConversation lumaUserConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class)
                .eq(LumaUserConversation::getUserId, user.getId())
                .eq(LumaUserConversation::getId, rid));
        if (lumaUserConversation == null) {
            throw BusinessRuntimeException.getInstance("视频不存在");
        }
        String url = LUMA_BASE_URL + "/luma/generations/" + lumaUserConversation.getTaskId();
        HttpRequest request = HttpRequest.get(url);
        request.header("Authorization", AUTH);

        String body = request.setConnectionTimeout(50000).execute().body();
        log.info("action body:{}", body);
        SubmitResult submitResult = Jsons.parseObject(body, SubmitResult.class);
        if (submitResult.getId() == null) {
            throw BusinessRuntimeException.getInstance("视频不存在");
        }
        if (FINISH_STATUS.contains(submitResult.getState())) {
            String id = submitResult.getId();
            if (!FINISH_STATUS.contains(lumaUserConversation.getStatus())) {
                lumaUserConversation.setStatus(submitResult.getState());
                SubmitResult.Video video = submitResult.getVideo();
                lumaUserConversation.setFinishTime(System.currentTimeMillis());
                //出水印
                String removeVideoWatermarkUrl = removeVideoWatermark(lumaUserConversation.getTaskId());
                lumaUserConversation.setVideoUrl(removeVideoWatermarkUrl);
                lumaUserConversation.setThumbnail(video.getThumbnail());
                TASK_EXECUTOR.execute(() -> {
                    //同步
                    sync(lumaUserConversation);
                });
            }
        }
        return submitResult;
    }

    @Override
    public void submitVideoNotify(String json) throws Exception {
        log.info("notifyHook:{}", json);
        if (StrUtil.isEmpty(json)) {
            return;
        }
        JSONObject jsons = JSONUtil.parseObj(json);
        VideoTask videoTask = JSONUtil.toBean(jsons, VideoTask.class);
        VideoTask.Data videoTaskData = videoTask.getData();
        String taskId = videoTaskData.getId();
        LumaUserConversation conversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class)
                .eq(LumaUserConversation::getTaskId, taskId)
                .last("limit 1"));
        if (conversation == null) {
            log.info("回调任务taskId:{}不存在", taskId);
            return;
        }
        conversation.setStatus(videoTaskData.getState());
        VideoTask.Video video = videoTaskData.getVideo();
        if (video != null) {
            String removeVideoWatermarkUrl = removeVideoWatermark(conversation.getTaskId());
            conversation.setVideoUrl(removeVideoWatermarkUrl);
            conversation.setThumbnail(video.getThumbnail());
        }
        conversation.setFailReason(videoTask.getFail_reason());
        if (videoTask.getFinish_time() != 0) {
            conversation.setFinishTime(videoTask.getFinish_time());
        }
        redisService.set(RedisService.key.LUMA_CONVERSATION.getName() + conversation.getTaskId(), conversation,RedisService.key.LUMA_CONVERSATION.getTimeout());

        TASK_EXECUTOR.execute(() -> {
            try {
                LumaUserConversation conversation1 = new LumaUserConversation();
                BeanUtil.copyProperties(conversation, conversation1);
                sync(conversation1);
            } catch (Exception e) {
                log.error("同步任务出错 taskId:{}, error:{}", conversation.getTaskId(), StringUtil.getErrorText(e));
            }
        });
    }

    @Override
    public SubmitResult extendedVideo(LumaUser user, SubmitVideoDTO videoDTO) throws Exception {
        String taskId = videoDTO.getTaskId();
        String prompt = videoDTO.getPrompt();
        Boolean expandPrompt = videoDTO.getExpandPrompt();
        String url = LUMA_BASE_URL + String.format("/luma/generations/%s/extend", taskId);

        LumaUserConversation lumaUserConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class)
                .eq(LumaUserConversation::getTaskId, taskId)
                .eq(LumaUserConversation::getUserId, user.getId())
                .last("limit 1"));
        if (lumaUserConversation == null) {
            throw BusinessRuntimeException.getInstance("视频不存在");
        }
        if (!StrUtil.equals("completed", lumaUserConversation.getStatus())) {
            throw BusinessRuntimeException.getInstance("视频未生成");
        }
        Map<String, Object> videoParam = MapUtil.builder(new HashMap<String, Object>())
                .put("aspect_ratio", "16:9")
                .build();
        videoParam.put("user_prompt", prompt);
        if (expandPrompt != null) {
            videoParam.put("expand_prompt", expandPrompt);
        }
        String picture = videoDTO.getPicture();
        String imageEndUrl = videoDTO.getImageEndUrl();
        if (StrUtil.isNotBlank(picture)) {
            videoParam.put("image_url", picture);
        }
        if (StrUtil.isNotBlank(imageEndUrl)) {
            videoParam.put("image_end_url", imageEndUrl);
        }
        //线上接口8085
        videoParam.put("notify_hook", lumaHost + (EnvCommonService.active_pre.equals(envCommonService.getEnv()) ? "/8083" : "/8085") + VIDEO_NOTIFY_URL);

        HttpRequest request = HttpRequest.post(url).setConnectionTimeout(10000).body(Jsons.toJson(videoParam));
        request.header("Authorization", AUTH);
        String body = request.execute().body();
        SubmitResult submitResult = Jsons.parseObject(body, SubmitResult.class);
        log.info("return body: {}", body);
        if (submitResult.getId() == null) {
            throw BusinessRuntimeException.getInstance("请输入正确的描述");
        }
        //重新设置taskId
        lumaUserConversation.setTaskId(submitResult.getId());
        lumaUserConversation.setStatus(submitResult.getState());
        lumaUserConversation.setPrompt(submitResult.getPrompt());
        conversationMapper.updateById(lumaUserConversation);
        return submitResult;
    }

    @Override
    public LumaUserConversation getConversationById(String taskId) {
        //从缓存中查询 该任务是否存在
        LumaUserConversation conversation;
        Object obj = redisService.get(RedisService.key.LUMA_CONVERSATION.getName() + taskId);
        if (obj == null) {
            conversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class)
                    .eq(LumaUserConversation::getTaskId, taskId)
                    .last("limit 1"));
            redisService.set(RedisService.key.LUMA_CONVERSATION.getName() + taskId, conversation, RedisService.key.LUMA_CONVERSATION.getTimeout());
        } else {
            try {
                conversation = (LumaUserConversation) obj;
            } catch (Exception e) {
                String jsonStr = JSONUtil.toJsonStr(obj);
                conversation = JSONUtil.parseObj(jsonStr).toBean(LumaUserConversation.class);
                redisService.set(RedisService.key.LUMA_CONVERSATION.getName() + taskId, conversation, RedisService.key.LUMA_CONVERSATION.getTimeout());
            }
        }
        return conversation;
    }


    @Override
    public List<LumaUserConversation> conversationListByIds(List<Long> ids, LumaUser user) {
        List<LumaUserConversation> list = new ArrayList<>();
        Integer selectCount = conversationMapper.selectCount(Wrappers.lambdaQuery(LumaUserConversation.class)
                .eq(LumaUserConversation::getUserId, user.getId())
                .in(LumaUserConversation::getTaskId, ids));
        if (selectCount != ids.size()) {
            throw BusinessRuntimeException.getInstance("视频不存在");
        }
        ids.forEach(id -> {
            String queryKey = RedisService.key.LUMA_QUERY.getName();
            Long count = redisService.incr(queryKey + id, 1L);
            if (count % 5 == 0) {
                try {
                    SubmitResult submitResult = getVideoTaskInfoByTaskId(user, id);
                    LumaUserConversation conversation = new LumaUserConversation();
                    SubmitResult.Video video = submitResult.getVideo();
                    conversation.setStatus(submitResult.getState());
                    conversation.setVideoUrl(video.getUrl());
                    if (!FINISH_STATUS.contains(conversation.getStatus())) {
                        conversation.setFinishTime(System.currentTimeMillis());
                    }
                    list.add(conversation);
                } catch (Exception e) {

                }
            } else {
                Object obj = redisService.get(RedisService.key.LUMA_CONVERSATION.getName() + id);
                LumaUserConversation conversation;
                try {
                    conversation = (LumaUserConversation) obj;
                } catch (Exception e) {
                    String jsonStr = JSONUtil.toJsonStr(obj);
                    conversation = JSONUtil.parseObj(jsonStr).toBean(LumaUserConversation.class);
                    redisService.set(RedisService.key.LUMA_CONVERSATION.getName() + id, conversation, RedisService.key.LUMA_CONVERSATION.getTimeout());
                }
                list.add(conversation);
            }
        });
        return list;
    }


    @Override
    public LumaUserConversation getVideoConversationByTid(LumaUser user, Long rid) {
        if (rid == null) {
            return null;
        }
        LumaUserConversation lumaUserConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class)
                .eq(LumaUserConversation::getUserId, user.getId())
                .in(LumaUserConversation::getId, rid)
                .select(LumaUserConversation::getId, LumaUserConversation::getTaskId, LumaUserConversation::getPrompt, LumaUserConversation::getVideoUrl, LumaUserConversation::getStartTime, LumaUserConversation::getFinishTime, LumaUserConversation::getTime, LumaUserConversation::getStatus, LumaUserConversation::getThumbnail));
        return lumaUserConversation;
    }

    @Override
    public String updateCover(Long id, Long userId) throws IOException {
        LumaUserConversation lumaUserConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class)
                .eq(LumaUserConversation::getId, id)
                .eq(LumaUserConversation::getUserId, userId));
        if (lumaUserConversation == null) {
            return null;
        }
        if (FINISH_STATUS.contains(lumaUserConversation.getStatus()) && StrUtil.isNotBlank(lumaUserConversation.getVideoUrl())) {
            String coverFrameByUrl = VideoCoverExtractor.extractVideoCoverFrameByUrl(lumaUserConversation.getVideoUrl());
            if (StrUtil.isNotBlank(coverFrameByUrl)) {
                conversationMapper.update(null, Wrappers.lambdaUpdate(LumaUserConversation.class)
                        .set(LumaUserConversation::getThumbnail, coverFrameByUrl)
                        .eq(LumaUserConversation::getId, lumaUserConversation.getId()));
                return coverFrameByUrl;
            }
        }
        return null;
    }

    public void sync(LumaUserConversation conversation) {
        if (FINISH_STATUS.contains(conversation.getStatus())) {
            //重复更新数据
            if (!redisService.setNx(conversation.getId().toString(), conversation.getId(), 10l)) {
                return;
            }
            log.info("同步任务：{},进度:{}",conversation.getTaskId(),conversation.getProgress());
            Long userId = conversation.getUserId();
            LumaUserConversation dbConversation = conversationMapper.selectOne(Wrappers.lambdaQuery(LumaUserConversation.class).eq(LumaUserConversation::getTaskId, conversation.getTaskId())
                    .eq(LumaUserConversation::getUserId, userId).orderByDesc(LumaUserConversation::getId).last("limit 1"));
            if (dbConversation != null) {
                if ("completed".equals(dbConversation.getStatus())) {
                    return;
                }
                conversation.setProgress("100%");
                conversation.setId(dbConversation.getId());
                conversation.setTime(Optional.ofNullable(dbConversation.getTime()).orElse(0) + 5);
                conversationMapper.updateById(conversation);
                TASK_EXECUTOR.execute(() -> {
                    try {
                        String videoUrl = conversation.getVideoUrl();
                        if (StringUtil.isNotBlank(videoUrl)) {
                            //上传视频至腾讯云
                            String f_videoUrl = uploadVideoUrl(videoUrl, "conversation" + dbConversation.getId(), conversation);
                            conversation.setVideoUrl(f_videoUrl);
                            conversationMapper.updateById(conversation);
                        }
                    } catch (IOException e) {
                        log.error("上传视频失败",e);
                    }
                });
                //失败返还次数
                if ("failed".equals(conversation.getStatus())){
                    LambdaUpdateWrapper<LumaUser> wrapper = Wrappers.lambdaUpdate(LumaUser.class)
                            .eq(LumaUser::getId, userId);
                    Object num = redisService.incr(RedisService.key.LUMA_NUM_LIMIT.getName() + userId, 1L);
                    wrapper.set(LumaUser::getNum, num);
                    lumaUserMapper.update(null, wrapper);
                }
            }
        }
    }


    private static String uploadVideoUrl(String url, String prefix, LumaUserConversation conversation) throws IOException {
        byte[] body;
        try {
            body = HttpUtil.downloadBytes(url);
        } catch (Exception e) {
            throw new IOException(e);
        }

        // 创建临时文件并将图像字节写入文件
        Path tempFile = Files.createTempFile(prefix, ".mp4");
        try (FileOutputStream fos = new FileOutputStream(tempFile.toFile())) {
            fos.write(body);
        }

        //抓取视频封面
        conversation.setThumbnail(VideoCoverExtractor.extractVideoCoverFrame(tempFile));

        String result = url;

        //上传视频并处理潜在的异常
        try (InputStream inputStream = Files.newInputStream(tempFile)) {
            String contentType = Files.probeContentType(tempFile);
            String originalFilename = tempFile.getFileName().toString();
            String ext = originalFilename.substring(originalFilename.lastIndexOf("."));
            result = COSUtil.upLoadVideo(inputStream, ext, contentType);
        } catch (Exception e) {
            log.error("上传视频异常:{}", StringUtil.getErrorText(e));
        } finally {
            try {
                Files.deleteIfExists(tempFile);
            } catch (IOException e) {
                log.error("删除临时文件异常:{}", StringUtil.getErrorText(e));
            }
        }

        log.info("上传视频结果：url:{},json:{}", url, result);
        return result;
    }

    /**
     * 去水印
     */
    public String removeVideoWatermark(String taskId) throws Exception {
        String url = LUMA_BASE_URL + String.format("/luma/generations/%s/download_video_url", taskId);
        HttpRequest request = HttpRequest.get(url).header("Authorization", AUTH).setConnectionTimeout(10000);
        HttpResponse execute = request.execute();
        String body = execute.body();
        JSONObject re = Jsons.parseObject(body, JSONObject.class);
        return re.getStr("url");
    }
}