EAI-Backend/src/main/java/com/njuzr/eaibackend/config/OssConfig.java

package com.njuzr.eaibackend.config;

import com.aliyun.oss.OSS;
import com.aliyun.oss.OSSClientBuilder;
import com.aliyun.oss.model.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author: Leonezhurui
 * @Date: 2024/3/13 - 09:15
 * @Package: EAI-Backend
 */

@Configuration
public class OssConfig {
    @Value("oss-cn-shanghai.aliyuncs.com")
    private String endpoint;

    @Value("${aliyun.oss.accessKeyId}")
    private String accessKeyId;

    @Value("${aliyun.oss.accessKeySecret}")
    private String accessKeySecret;

    @Value("${aliyun.oss.bucketName}")
    private String bucketName;

    private OSS ossClient;


    // 存储uploadId的映射，key为fileMd5
    private final ConcurrentHashMap<String, String> uploadIdMap = new ConcurrentHashMap<>();
    // 存储已上传分片的映射，key为fileMd5
    private final ConcurrentHashMap<String, List<PartETag>> partETagsMap = new ConcurrentHashMap<>();
    @Bean
    public OSS ossClient() {
        if(null == this.ossClient){
            this.ossClient = new OSSClientBuilder().build(endpoint, accessKeyId, accessKeySecret);
        }
         return this.ossClient;
    }

    /**
     * 初始化分片上传
     */
    public String initMultipartUpload(String fileMd5, String fileName) {
        String objectName = "uploads/" + fileMd5+"-" + fileName;
        InitiateMultipartUploadRequest request = new InitiateMultipartUploadRequest(bucketName, objectName);
        InitiateMultipartUploadResult result = ossClient.initiateMultipartUpload(request);
        String uploadId = result.getUploadId();
        System.out.println("init"+ uploadId);
        uploadIdMap.put(fileMd5, uploadId);
        partETagsMap.put(fileMd5, new ArrayList<>());
        return uploadId;
    }

    public boolean isUploadInitialized(String fileMd5) {
        return uploadIdMap.containsKey(fileMd5);
    }

    /**
     * 上传分片
     */
    public boolean uploadPart(String fileName,String fileMd5, int partNumber, InputStream inputStream, long partSize) throws IOException {
        String uploadId = uploadIdMap.get(fileMd5);
        if (uploadId == null) {
            throw new RuntimeException("未找到对应的uploadId，请先初始化分片上传");
        }
        String objectName = "uploads/" + fileMd5 + "-" + fileName;
        UploadPartRequest uploadPartRequest = new UploadPartRequest();
        uploadPartRequest.setBucketName(bucketName);
        uploadPartRequest.setKey(objectName);
        uploadPartRequest.setUploadId(uploadId);
        uploadPartRequest.setInputStream(inputStream);
        uploadPartRequest.setPartSize(partSize);
        uploadPartRequest.setPartNumber(partNumber);

        UploadPartResult uploadPartResult = ossClient.uploadPart(uploadPartRequest);
        partETagsMap.get(fileMd5).add(uploadPartResult.getPartETag());
        return true;
    }

    /**
     * 完成分片上传
     */
    public String completeMultipartUpload(String fileMd5, String fileName) {
        String uploadId = uploadIdMap.get(fileMd5);
        if (uploadId == null) {
            throw new RuntimeException("未找到对应的uploadId");
        }
        String objectName = "uploads/" + fileMd5 + "-" + fileName;
        CompleteMultipartUploadRequest completeRequest = new CompleteMultipartUploadRequest(
                bucketName,
                objectName,
                uploadId,
                partETagsMap.get(fileMd5));

        CompleteMultipartUploadResult result = ossClient.completeMultipartUpload(completeRequest);
        // 清理资源
        uploadIdMap.remove(fileMd5);
        partETagsMap.remove(fileMd5);
        // 生成文件访问URL
        return generateFileUrl(objectName);
    }

    private String generateFileUrl(String objectName) {
        return "https://" + bucketName + "." + endpoint.replace("https://", "") + "/" + objectName;
    }

    /**
     * 取消分片上传
     */
    public void abortMultipartUpload(String fileMd5, String filName) {
        String uploadId = uploadIdMap.get(fileMd5);
        if (uploadId != null) {
            String objectName = "uploads/" + fileMd5 + "-" + filName;
            AbortMultipartUploadRequest abortRequest =
                    new AbortMultipartUploadRequest(bucketName, objectName, uploadId);
            ossClient.abortMultipartUpload(abortRequest);

            uploadIdMap.remove(fileMd5);
            partETagsMap.remove(fileMd5);
        }
    }

    /**
     * 获取已上传的分片数量
     */
    public int getChunkCount(String fileMd5) {
        List<PartETag> partETags = partETagsMap.get(fileMd5);
        return partETags != null ? partETags.size() : 0;
    }

    /**
     * 关闭OSS客户端
     */
    public void shutdown() {
        if (ossClient != null) {
            ossClient.shutdown();
        }
    }
}