vdi/nex-be/src/main/java/com/unisinsight/project/controller/FileChunkController.java

package com.unisinsight.project.controller;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.unisinsight.project.entity.dao.Image;
import com.unisinsight.project.exception.BaseErrorCode;
import com.unisinsight.project.exception.Result;
import com.unisinsight.project.mapper.ImageMapper;
import com.unisinsight.project.util.DigestUtil;
import io.swagger.annotations.*;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.ObjectUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpMethod;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.Resource;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 大文件分片上传控制器
 */
@RestController
@RequestMapping("/api/files")
@Api(tags = "文件分片上传接口")
@Slf4j
public class FileChunkController {

    // 临时目录，用于存储上传的分片
    @Value("${file.upload.temp-dir:${java.io.tmpdir}/chunked-uploads}")
    private String tempDir;

    // 最终文件存储目录
    @Value("${file.upload.dir:${user.home}/uploads}")
    private String uploadDir;

    // 请求bt配置
    @Value("${file.upload.bt-url}")
    private String btUrl;

    @Resource
    private RestTemplate restTemplate;


    // 存储每个文件的分片信息
    private final Map<String, FileUploadInfo> fileUploadMap = new ConcurrentHashMap<>();

    @Resource
    private ImageMapper imageMapper;

    /**
     * 上传文件分片
     *
     * @param chunk       分片文件
     * @param fileId      文件唯一标识符
     * @param chunkNumber 当前分片编号(从1开始)
     * @param totalChunks 总分片数
     * @param fileName    原始文件名
     * @param totalSize   文件总大小
     * @return 上传结果
     */
    @PostMapping("/upload-chunk")
    @ApiOperation(value = "上传文件分片", notes = "上传单个文件分片，当所有分片上传完成后自动合并文件")
    @ApiImplicitParams({
            @ApiImplicitParam(name = "chunk", value = "文件分片", required = true, dataType = "__File", paramType = "form"),
            @ApiImplicitParam(name = "chunk_size", value = "文件分片大小", required = true, dataType = "int", paramType = "query"),
            @ApiImplicitParam(name = "chunk_md5", value = "文件分片md5", required = true, dataType = "String", paramType = "query"),
            @ApiImplicitParam(name = "file_id", value = "文件唯一标识符", required = true, dataType = "String", paramType = "query"),
            @ApiImplicitParam(name = "shard_index", value = "当前分片编号(从1开始)", required = true, dataType = "int", paramType = "query"),
            @ApiImplicitParam(name = "shard_total", value = "总分片数", required = true, dataType = "int", paramType = "query"),
            @ApiImplicitParam(name = "file_name", value = "原始文件名", required = true, dataType = "String", paramType = "query"),
            @ApiImplicitParam(name = "file_size", value = "文件总大小", required = true, dataType = "long", paramType = "query")
    })
    @ApiResponses({
            @ApiResponse(code = 200, message = "上传成功"),
            @ApiResponse(code = 500, message = "服务器内部错误")
    })
    public ResponseEntity<Map<String, Object>> uploadChunk(
            @RequestParam("chunk") MultipartFile chunk,
            @RequestParam("chunk_size") int chunkSize,
            @RequestParam("chunk_md5") String chunkMd5,
            @RequestParam("file_id") String fileId,
            @RequestParam("shard_index") int chunkNumber,
            @RequestParam("shard_total") int totalChunks,
            @RequestParam("file_name") String fileName,
            @RequestParam("file_size") long totalSize) {

        Map<String, Object> response = new HashMap<>();

        try {
            String md5 = DigestUtil.encryptMd5(chunk.getBytes());
            if (!chunkMd5.equals(md5)) {
                log.info("分片文件md5校验失败,chunkMd5:{},md5:{}", chunkMd5, md5);
                throw new RuntimeException("分片文件md5校验失败");
            }
            QueryWrapper<Image> wrapper = new QueryWrapper<>();
            wrapper.lambda().eq(Image::getImageName, fileName);
            List<Image> imageList = imageMapper.selectList(wrapper);
            if (CollectionUtil.isNotEmpty(imageList)) {
                response.put("success", false);
                response.put("status", "error");
                response.put("message", "当前文件已经上传");
                return ResponseEntity.status(200).body(response);
            }
            // 创建临时目录
            Path fileTempDir = Paths.get(tempDir, fileId);
            if (!Files.exists(fileTempDir)) {
                Files.createDirectories(fileTempDir);
            }
            // 保存分片文件
            String chunkFileName = String.format("%05d.part", chunkNumber);
            log.info("保存分片文件: {}", chunkFileName);
            Path chunkFilePath = fileTempDir.resolve(chunkFileName);
            chunk.transferTo(chunkFilePath);

            // 更新文件上传信息
            FileUploadInfo uploadInfo = fileUploadMap.computeIfAbsent(fileId,
                    id -> new FileUploadInfo(id, fileName, totalChunks, totalSize));
            uploadInfo.addUploadedChunk(chunkNumber);
//            log.info("更新文件上传信息: {}", JSONUtil.toJsonStr(uploadInfo));
            // 检查是否所有分片都已上传
            if (uploadInfo.isUploadComplete()) {
                // 合并文件
                Path finalDir = Paths.get(uploadDir);
                if (!Files.exists(finalDir)) {
                    Files.createDirectories(finalDir);
                }
                Path finalFilePath = finalDir.resolve(fileName);
                log.info("合并所有分片文件: {}", finalFilePath);
                mergeChunks(fileId, finalFilePath, totalChunks);

                // 清理临时文件
                log.info("清理临时文件: {}", fileId);
                cleanupTempFiles(fileId);

                // 从上传映射中移除
                fileUploadMap.remove(fileId);

                Image image = new Image();
                image.setImageName(fileName);
                image.setStoragePath(String.valueOf(finalFilePath));
                image.setImageStatus(1);
                int insert = imageMapper.insert(image);
                log.info("镜像新增insert:{}", insert);
                if (insert == 1) {

                    // 异步执行创建和做种操作
                    CompletableFuture.runAsync(() -> {
                        try {
                            String url = btUrl + "/vdi/start?sourceFile=%s&torrentFile=%s";
                            url = String.format(url, finalFilePath, finalFilePath + ".torrent");
                            log.info("请求bt创建接口参数: {}", url);
                            ResponseEntity<Boolean> responseEntity = restTemplate.exchange(url, HttpMethod.GET, null, Boolean.class);
                            log.info("请求bt创建接口返回: {}", JSONUtil.toJsonStr(responseEntity));
                            HttpStatus statusCode = responseEntity.getStatusCode();
                            if (statusCode != HttpStatus.OK) {
                                boolean result = Boolean.TRUE.equals(responseEntity.getBody());
                                if (result) {
                                    log.info("请求bt创建接口成功");
                                    QueryWrapper<Image> imageQueryWrapper = new QueryWrapper<>();
                                    imageQueryWrapper.lambda().eq(Image::getImageName, fileName);
                                    Image imageBt = imageMapper.selectOne(imageQueryWrapper);
                                    if (ObjectUtils.isNotEmpty(imageBt)) {
                                        imageBt.setBtPath(finalFilePath + ".torrent");
                                        int update = imageMapper.updateById(imageBt);
                                        log.info("镜像bt更新insert:{}", update);
                                    } else {
                                        log.info("镜像bt更新查询失败:{}", fileName);
                                    }
                                }
                            }
                        } catch (Exception e) {
                            log.error("请求bt创建接口失败: {}", e.getMessage(), e);
                        }
                    });
                    response.put("status", "completed");
                    response.put("message", "文件上传并合并完成");
                    response.put("filePath", finalFilePath.toString());
                } else {
                    throw new RuntimeException("文件上传失败");
                }

            } else {
                response.put("status", "uploading");
                response.put("message", "分片上传成功");
                response.put("uploadedChunks", uploadInfo.getUploadedChunks().size());
                response.put("totalChunks", totalChunks);
            }

            response.put("success", true);
            return ResponseEntity.ok(response);

        } catch (Exception e) {
            response.put("success", false);
            response.put("status", "error");
            response.put("message", "上传失败");
            log.info("上次失败清理临时文件: {},error:{}", fileId, e.getMessage(), e);
            try {
                cleanupTempFiles(fileId);
                cleanUploadFile(fileName);
            } catch (IOException ex) {
                log.error("清理临时文件失败,fileId:{}, {}", fileId, ex.getMessage(), ex);
            }
            return ResponseEntity.status(500).body(response);
        }
    }

    @PostMapping("/cancel/upload")
    @ApiOperation(value = "取消上传文件")
    public Result<?> cancelUpload(@RequestParam("file_id") String fileId) {
        if (ObjectUtils.isEmpty(fileId)) {
            return Result.errorResult(BaseErrorCode.PARAMETERS_EMPTY);
        }
        log.info("取消上传,清理临时文件: {}", fileId);
        try {
            // 从上传映射中移除
            fileUploadMap.remove(fileId);
            cleanupTempFiles(fileId);
        } catch (IOException ex) {
            log.error("清理临时文件失败,fileId:{}, {}", fileId, ex.getMessage(), ex);
        }
        return Result.successResult();
    }

    /**
     * 查询文件上传状态
     *
     * @param fileId 文件唯一标识符
     * @return 上传状态信息
     */
    @GetMapping("/upload-status/{fileId}")
    @ApiOperation("查询文件上传状态")
    public ResponseEntity<Map<String, Object>> getUploadStatus(@PathVariable String fileId) {
        Map<String, Object> response = new HashMap<>();

        FileUploadInfo uploadInfo = fileUploadMap.get(fileId);
        if (uploadInfo == null) {
            // 检查文件是否已经完成上传并合并
            try {
                Path finalFilePath = Paths.get(uploadDir, fileId);
                if (Files.exists(finalFilePath)) {
                    response.put("status", "completed");
                    response.put("message", "文件上传已完成");
                    response.put("filePath", finalFilePath.toString());
                } else {
                    response.put("status", "not_found");
                    response.put("message", "文件上传信息不存在");
                }
            } catch (Exception e) {
                response.put("status", "error");
                response.put("message", "查询状态失败: " + e.getMessage());
            }
        } else {
            response.put("status", "uploading");
            response.put("uploadedChunks", uploadInfo.getUploadedChunks().size());
            response.put("totalChunks", uploadInfo.getTotalChunks());
            response.put("progress", (double) uploadInfo.getUploadedChunks().size() / uploadInfo.getTotalChunks());
        }

        response.put("success", true);
        return ResponseEntity.ok(response);
    }

    /**
     * 合并所有分片文件
     *
     * @param fileId      文件唯一标识符
     * @param outputPath  合并后的文件路径
     * @param totalChunks 总分片数
     * @throws IOException IO异常
     */
    private void mergeChunks(String fileId, Path outputPath, int totalChunks) throws IOException {
        try (OutputStream outputStream = Files.newOutputStream(outputPath)) {
            Path fileTempDir = Paths.get(tempDir, fileId);

            // 按顺序合并分片
            for (int i = 1; i <= totalChunks; i++) {
                String chunkFileName = String.format("%05d.part", i);
                Path chunkPath = fileTempDir.resolve(chunkFileName);

                if (!Files.exists(chunkPath)) {
                    throw new IOException("缺少分片文件: " + chunkFileName);
                }

                // 将分片内容追加到输出文件
                Files.copy(chunkPath, outputStream);
            }
        }
    }

    /**
     * 清理临时分片文件
     *
     * @param fileId 文件唯一标识符
     * @throws IOException IO异常
     */
    private void cleanupTempFiles(String fileId) throws IOException {
        Path fileTempDir = Paths.get(tempDir, fileId);
        if (Files.exists(fileTempDir)) {
            // 递归删除临时目录及其内容
            Files.walk(fileTempDir)
                    .sorted(Comparator.reverseOrder())
                    .map(Path::toFile)
                    .forEach(File::delete);
        }
    }

    /**
     * 根据文件名删除已上传的文件
     *
     * @param fileName 文件名
     * @throws IOException IO异常
     */
    private void cleanUploadFile(String fileName) throws IOException {
        Path filePath = Paths.get(uploadDir, fileName);
        if (Files.exists(filePath)) {
            // 删除文件
            Files.delete(filePath);
            log.info("已删除文件: {}", filePath);
        } else {
            log.warn("文件不存在，无需删除: {}", filePath);
        }
    }


    /**
     * 文件上传信息类
     */
    @Data
    private static class FileUploadInfo {
        private final String fileId;
        private final String fileName;
        private final int totalChunks;
        private final long totalSize;
        private final Set<Integer> uploadedChunks;

        public FileUploadInfo(String fileId, String fileName, int totalChunks, long totalSize) {
            this.fileId = fileId;
            this.fileName = fileName;
            this.totalChunks = totalChunks;
            this.totalSize = totalSize;
            this.uploadedChunks = ConcurrentHashMap.newKeySet();
        }

        public void addUploadedChunk(int chunkNumber) {
            uploadedChunks.add(chunkNumber);
        }

        public boolean isUploadComplete() {
            return uploadedChunks.size() == totalChunks;
        }

    }

}