stray-animals-back/src/main/java/com/youlai/boot/mini/controller/AiGenerationController.java

package com.youlai.boot.mini.controller;

import com.youlai.boot.common.annotation.Log;
import com.youlai.boot.common.annotation.RepeatSubmit;
import com.youlai.boot.common.enums.ActionTypeEnum;
import com.youlai.boot.common.enums.LogModuleEnum;
import com.youlai.boot.common.result.Result;
import com.youlai.boot.framework.security.util.SecurityUtils;
import com.youlai.boot.mini.model.form.AiFourPanelGenerateForm;
import com.youlai.boot.mini.model.form.AiSingleImageGenerateForm;
import com.youlai.boot.mini.model.form.AiVideoGenerateForm;
import com.youlai.boot.mini.model.vo.AiTaskCallbackVO;
import com.youlai.boot.mini.model.vo.AiVideoCallbackVO;
import com.youlai.boot.mini.service.AiGenerationService;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.enums.ParameterIn;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.validation.Valid;
import lombok.RequiredArgsConstructor;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.Collections;
import java.util.List;

@Tag(name = "AI生成图片视频相关接口")
@RestController
@RequestMapping("/api/v1/mini/ai/generation")
@RequiredArgsConstructor
@Valid
public class AiGenerationController {

    private final AiGenerationService aiGenerationService;

    @Operation(summary = "上传AI生成参考文件", operationId = "AiReferenceSaveFile")
    @PostMapping(value = "/upload-reference", consumes = MediaType.MULTIPART_FORM_DATA_VALUE)
    @RepeatSubmit
    @Log(module = LogModuleEnum.AI_TASK_MEDIA, value = ActionTypeEnum.INSERT)
    public Result<List<String>> uploadReferenceFile(
            @RequestPart(name = "images", required = false) List<MultipartFile> images,
            @RequestPart(name = "videos", required = false) List<MultipartFile> videos
    ) {
        Long userId = SecurityUtils.getUserId();
        List<String> urlList = aiGenerationService.uploadReferenceFile(images, videos, userId);
        return Result.success(urlList);
    }

    @Operation(summary = "提交单图生成任务")
    @PostMapping("/generate-single-image")
    @Log(module = LogModuleEnum.AI_GENERATION_TASK, value = ActionTypeEnum.INSERT)
    public Result<String> generateSingleImage(@Valid @RequestBody AiSingleImageGenerateForm form) {
        Long userId = SecurityUtils.getUserId();
        String taskUuid = aiGenerationService.createAndGenerateImage(form, userId);
        return Result.success(taskUuid);
    }

    //TODO: 后续增加用户小程序订阅通知
    @Operation(summary = "单图任务回调接口")
    @PostMapping("/single-image/task/callback")
    @Log(module = LogModuleEnum.AI_GENERATION_TASK, value = ActionTypeEnum.UPDATE)
    public Result<Boolean> taskCallback(@Valid @RequestBody AiTaskCallbackVO vo) {
        boolean success = aiGenerationService.handleTaskCallback(vo);
        return Result.success(success);
    }

    @Operation(summary = "提交四宫格漫画生成任务")
    @PostMapping("/generate-four-panel")
    @Log(module = LogModuleEnum.AI_GENERATION_TASK, value = ActionTypeEnum.INSERT)
    public Result<String> generateFourPanelImage(@Valid @RequestBody AiFourPanelGenerateForm form) {
        Long userId = SecurityUtils.getUserId();
        String taskUuid = aiGenerationService.createAndGenerateFourPanel(form, userId);
        return Result.success(taskUuid);
    }

    @Operation(summary = "提交视频生成任务")
    @PostMapping("/generate-video")
    @Log(module = LogModuleEnum.AI_GENERATION_TASK, value = ActionTypeEnum.INSERT)
    public Result<String> generateVideo(@Valid @RequestBody AiVideoGenerateForm form) {
        Long userId = SecurityUtils.getUserId();
        String taskUuid = aiGenerationService.createAndGenerateVideo(form, userId);
        return Result.success(taskUuid);
    }

    //TODO: 后续增加用户小程序订阅通知
    @Operation(summary = "视频任务回调接口")
    @PostMapping("/video/task/callback")
    @Log(module = LogModuleEnum.AI_GENERATION_TASK, value = ActionTypeEnum.UPDATE)
    public Result<Boolean> videoTaskCallback(@Valid @RequestBody AiVideoCallbackVO vo) {
        boolean success = aiGenerationService.handleVideoTaskCallback(vo);
        return Result.success(success);
    }

}