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

package com.unisinsight.project.controller;

import io.swagger.annotations.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

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.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

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

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

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

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

    /**
     * 上传文件分片
     * 
     * @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 = "fileId", value = "文件唯一标识符", required = true, dataType = "String", paramType = "query"),
            @ApiImplicitParam(name = "chunkNumber", value = "当前分片编号(从1开始)", required = true, dataType = "int", paramType = "query"),
            @ApiImplicitParam(name = "totalChunks", value = "总分片数", required = true, dataType = "int", paramType = "query"),
            @ApiImplicitParam(name = "fileName", value = "原始文件名", required = true, dataType = "String", paramType = "query"),
            @ApiImplicitParam(name = "totalSize", 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("fileId") String fileId,
            @RequestParam("chunkNumber") int chunkNumber,
            @RequestParam("totalChunks") int totalChunks,
            @RequestParam("fileName") String fileName,
            @RequestParam("totalSize") long totalSize) {
        
        Map<String, Object> response = new HashMap<>();
        
        try {
            // 创建临时目录
            Path fileTempDir = Paths.get(tempDir, fileId);
            if (!Files.exists(fileTempDir)) {
                Files.createDirectories(fileTempDir);
            }

            // 保存分片文件
            String chunkFileName = String.format("%05d.part", chunkNumber);
            Path chunkFilePath = fileTempDir.resolve(chunkFileName);
            chunk.transferTo(chunkFilePath);

            // 更新文件上传信息
            FileUploadInfo uploadInfo = fileUploadMap.computeIfAbsent(fileId, 
                id -> new FileUploadInfo(id, fileName, totalChunks, totalSize));
            uploadInfo.addUploadedChunk(chunkNumber);

            // 检查是否所有分片都已上传
            if (uploadInfo.isUploadComplete()) {
                // 合并文件
                Path finalDir = Paths.get(uploadDir);
                if (!Files.exists(finalDir)) {
                    Files.createDirectories(finalDir);
                }

                Path finalFilePath = finalDir.resolve(fileName);
                mergeChunks(fileId, finalFilePath, totalChunks);

                // 清理临时文件
                cleanupTempFiles(fileId);
                
                // 从上传映射中移除
                fileUploadMap.remove(fileId);

                response.put("status", "completed");
                response.put("message", "文件上传并合并完成");
                response.put("filePath", finalFilePath.toString());
            } 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("message", "上传失败: " + e.getMessage());
            return ResponseEntity.status(500).body(response);
        }
    }

    /**
     * 查询文件上传状态
     * 
     * @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);
        }
    }

    /**
     * 文件上传信息类
     */
    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;
        }

        public String getFileId() {
            return fileId;
        }

        public String getFileName() {
            return fileName;
        }

        public int getTotalChunks() {
            return totalChunks;
        }

        public long getTotalSize() {
            return totalSize;
        }

        public Set<Integer> getUploadedChunks() {
            return uploadedChunks;
        }
    }

    @GetMapping("/test")
    @ApiOperation("测试")
    public ResponseEntity<String> getUploadStatus() {
       return ResponseEntity.ok("ok");
    }
}