SeeyonFileSystem/server/service/chunk_upload_service.go

package service

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"

	"seeyon-filesystem/config"
	"seeyon-filesystem/dto"
	"seeyon-filesystem/model"
	"seeyon-filesystem/repository"
	"seeyon-filesystem/storage"
	"seeyon-filesystem/utils"

	"github.com/google/uuid"
)

// ChunkUploadService 分片上传服务
type ChunkUploadService struct {
	fileRepo *repository.FileRepository
}

func NewChunkUploadService() *ChunkUploadService {
	return &ChunkUploadService{
		fileRepo: repository.NewFileRepository(),
	}
}

// InitUpload 初始化分片上传
// 根据存储策略自动选择: 本地模式(服务端中转) 或 MinIO模式(客户端直传)
func (s *ChunkUploadService) InitUpload(fileName string, fileSize int64, md5 string, folderID, ownerID, policyID uint64) (*dto.ChunkUploadInitVO, error) {
	// 秒传检查
	if md5 != "" {
		existing, _ := s.fileRepo.FindByMD5(md5)
		if existing != nil {
			return &dto.ChunkUploadInitVO{
				UploadID:  "",
				Instant:   true,
				FileID:    existing.ID,
				ChunkSize: 0,
				Total:     0,
			}, nil
		}
	}

	// 计算分片参数
	chunkSize := int64(5 * 1024 * 1024)
	if fileSize < chunkSize {
		chunkSize = fileSize
	}
	totalChunks := int((fileSize + chunkSize - 1) / chunkSize)

	uploadID := uuid.New().String()

	// 生成存储路径
	ext := utils.GetExtension(fileName)
	var storageKey string
	if md5 != "" {
		storageKey = storage.GenerateStorageKey(md5, ext)
	} else {
		storageKey = storage.GenerateStorageKey(fmt.Sprintf("%d-%d", ownerID, time.Now().UnixNano()), ext)
	}

	// 创建上传会话
	tempDir := filepath.Join(config.AppConfig.Storage.LocalBasePath, ".tmp", uploadID)
	os.MkdirAll(tempDir, 0755)

	session := &model.UploadSession{
		ID:          uploadID,
		FileName:    fileName,
		FileSize:    fileSize,
		ChunkSize:   chunkSize,
		TotalChunks: totalChunks,
		MD5:         md5,
		FolderID:    folderID,
		OwnerID:     ownerID,
		PolicyID:    policyID,
		Status:      1,
		TempDir:     storageKey, // 复用字段存储目标storageKey
	}
	config.DB.Create(session)

	// 检查是否为MinIO存储
	engine, err := storage.GetEngineByPolicyID(policyID)
	if err == nil {
		if minioEngine, ok := engine.(*storage.MinIOEngine); ok {
			return s.initMinIOUpload(minioEngine, session, storageKey)
		}
	}

	// 本地存储模式
	return &dto.ChunkUploadInitVO{
		UploadID:  uploadID,
		Instant:   false,
		ChunkSize: chunkSize,
		Total:     totalChunks,
		Mode:      "local",
	}, nil
}

// initMinIOUpload 初始化MinIO直传模式
func (s *ChunkUploadService) initMinIOUpload(minioEngine *storage.MinIOEngine, session *model.UploadSession, storageKey string) (*dto.ChunkUploadInitVO, error) {
	// 为每个分片生成预签名上传URL
	parts := make([]dto.ChunkPresignURL, session.TotalChunks)
	for i := 0; i < session.TotalChunks; i++ {
		chunkKey := fmt.Sprintf("%s.chunk_%06d", storageKey, i)
		presignedURL, err := minioEngine.GetPresignedChunkURL(chunkKey, 24*time.Hour)
		if err != nil {
			presignedURL = ""
		}
		parts[i] = dto.ChunkPresignURL{Index: i, URL: presignedURL}
	}

	return &dto.ChunkUploadInitVO{
		UploadID:  session.ID,
		ChunkSize: session.ChunkSize,
		Total:     session.TotalChunks,
		Mode:      "minio",
		Parts:     parts,
	}, nil
}

// UploadChunk 上传单个分片(本地模式)
func (s *ChunkUploadService) UploadChunk(uploadID string, index int, reader io.Reader, checksum string) (*dto.ChunkUploadChunkVO, error) {
	var session model.UploadSession
	if err := config.DB.First(&session, uploadID).Error; err != nil {
		return nil, errors.New("上传会话不存在")
	}
	if session.Status != 1 {
		return nil, errors.New("上传会话已结束")
	}
	if index < 0 || index >= session.TotalChunks {
		return nil, errors.New("分片索引无效")
	}

	uploaded := parseChunkIndices(session.UploadedChunks)
	if uploaded[index] {
		return &dto.ChunkUploadChunkVO{Index: index, Status: "exists"}, nil
	}

	chunkPath := filepath.Join(config.AppConfig.Storage.LocalBasePath, ".tmp", uploadID, fmt.Sprintf("chunk_%06d", index))
	dst, err := os.Create(chunkPath)
	if err != nil {
		return nil, fmt.Errorf("创建分片文件失败: %w", err)
	}

	hash := sha256.New()
	writer := io.MultiWriter(dst, hash)
	size, err := io.Copy(writer, reader)
	dst.Close()
	if err != nil {
		os.Remove(chunkPath)
		return nil, fmt.Errorf("写入分片失败: %w", err)
	}

	actualChecksum := hex.EncodeToString(hash.Sum(nil))
	if checksum != "" && checksum != actualChecksum {
		os.Remove(chunkPath)
		return nil, fmt.Errorf("分片校验失败: 期望 %s, 实际 %s", checksum, actualChecksum)
	}

	uploaded[index] = true
	chunkStr := buildChunkIndices(uploaded)
	config.DB.Model(&session).Update("uploaded_chunks", chunkStr)

	return &dto.ChunkUploadChunkVO{Index: index, Size: size, Status: "ok"}, nil
}

// UploadChunkMinIO 确认MinIO分片上传完成(客户端直传后调用)
func (s *ChunkUploadService) UploadChunkMinIO(uploadID string, index int, etag string) (*dto.ChunkUploadChunkVO, error) {
	var session model.UploadSession
	if err := config.DB.First(&session, uploadID).Error; err != nil {
		return nil, errors.New("上传会话不存在")
	}
	if session.Status != 1 {
		return nil, errors.New("上传会话已结束")
	}

	uploaded := parseChunkIndices(session.UploadedChunks)
	uploaded[index] = true
	chunkStr := buildChunkIndices(uploaded)
	config.DB.Model(&session).Update("uploaded_chunks", chunkStr)

	return &dto.ChunkUploadChunkVO{Index: index, Status: "ok", ETag: etag}, nil
}

// MergeChunks 合并分片完成上传
func (s *ChunkUploadService) MergeChunks(uploadID string) (*dto.FileVO, error) {
	var session model.UploadSession
	if err := config.DB.First(&session, uploadID).Error; err != nil {
		return nil, errors.New("上传会话不存在")
	}
	if session.Status != 1 {
		return nil, errors.New("上传会话已结束")
	}

	uploaded := parseChunkIndices(session.UploadedChunks)
	for i := 0; i < session.TotalChunks; i++ {
		if !uploaded[i] {
			return nil, fmt.Errorf("分片 %d 未上传", i)
		}
	}

	storageKey := session.TempDir // 目标storageKey

	// 检测MinIO模式
	engine, _ := storage.GetEngineByPolicyID(session.PolicyID)
	if minioEngine, ok := engine.(*storage.MinIOEngine); ok {
		return s.mergeMinIO(minioEngine, &session, storageKey)
	}

	// 本地模式: 合并分片文件
	return s.mergeLocal(&session, storageKey)
}

// mergeLocal 本地模式合并
func (s *ChunkUploadService) mergeLocal(session *model.UploadSession, storageKey string) (*dto.FileVO, error) {
	tmpDir := filepath.Join(config.AppConfig.Storage.LocalBasePath, ".tmp", session.ID)
	mergedPath := filepath.Join(tmpDir, "merged")

	merged, err := os.Create(mergedPath)
	if err != nil {
		return nil, fmt.Errorf("创建合并文件失败: %w", err)
	}

	for i := 0; i < session.TotalChunks; i++ {
		chunkPath := filepath.Join(tmpDir, fmt.Sprintf("chunk_%06d", i))
		chunk, err := os.Open(chunkPath)
		if err != nil {
			merged.Close()
			return nil, fmt.Errorf("读取分片 %d 失败: %w", i, err)
		}
		io.Copy(merged, chunk)
		chunk.Close()
	}
	merged.Close()

	fileMD5, _ := utils.FileMD5(mergedPath)
	fileSHA256, _ := utils.FileSHA256(mergedPath)

	if session.MD5 != "" && session.MD5 != fileMD5 {
		os.RemoveAll(tmpDir)
		return nil, fmt.Errorf("MD5校验失败: 期望 %s, 实际 %s", session.MD5, fileMD5)
	}

	// 秒传检查
	existing, _ := s.fileRepo.FindByMD5(fileMD5)
	if existing != nil {
		os.RemoveAll(tmpDir)
		newFile := &model.File{
			ID: utils.GenID(), Name: session.FileName, Extension: utils.GetExtension(session.FileName),
			FolderID: session.FolderID, OwnerID: session.OwnerID, Size: session.FileSize,
			MD5: fileMD5, StorageKey: existing.StorageKey, StoragePolicyID: existing.StoragePolicyID,
			MimeType: utils.GetMimeType(utils.GetExtension(session.FileName)), Status: 1,
		}
		s.fileRepo.Create(newFile)
		return fileToVO(newFile), nil
	}

	// 上传到存储引擎
	ext := utils.GetExtension(session.FileName)
	engine, _ := storage.GetEngineByPolicyID(session.PolicyID)
	reader, _ := os.Open(mergedPath)
	defer reader.Close()
	if err := engine.Upload(storageKey, reader, session.FileSize); err != nil {
		os.RemoveAll(tmpDir)
		return nil, fmt.Errorf("存储文件失败: %w", err)
	}

	newFile := &model.File{
		ID: utils.GenID(), Name: session.FileName, Extension: ext, FolderID: session.FolderID, OwnerID: session.OwnerID,
		Size: session.FileSize, MD5: fileMD5, SHA256: fileSHA256, StorageKey: storageKey,
		StoragePolicyID: session.PolicyID, MimeType: utils.GetMimeType(ext), Status: 1,
	}
	if err := s.fileRepo.Create(newFile); err != nil {
		engine.Delete(storageKey)
		os.RemoveAll(tmpDir)
		return nil, fmt.Errorf("保存文件记录失败: %w", err)
	}

	config.DB.Model(session).Updates(map[string]interface{}{"status": 2, "file_id": newFile.ID, "md5": fileMD5, "sha256": fileSHA256})
	os.RemoveAll(tmpDir)
	return fileToVO(newFile), nil
}

// mergeMinIO MinIO模式合并
func (s *ChunkUploadService) mergeMinIO(minioEngine *storage.MinIOEngine, session *model.UploadSession, storageKey string) (*dto.FileVO, error) {
	// 构建分片key列表
	chunkKeys := make([]string, session.TotalChunks)
	for i := 0; i < session.TotalChunks; i++ {
		chunkKeys[i] = fmt.Sprintf("%s.chunk_%06d", storageKey, i)
	}

	// 合并分片
	if err := minioEngine.MergeChunks(chunkKeys, storageKey); err != nil {
		return nil, fmt.Errorf("MinIO合并失败: %w", err)
	}

	ext := utils.GetExtension(session.FileName)
	fileMD5 := session.MD5 // 使用客户端提供的MD5

	newFile := &model.File{
		ID: utils.GenID(), Name: session.FileName, Extension: ext, FolderID: session.FolderID, OwnerID: session.OwnerID,
		Size: session.FileSize, MD5: fileMD5, StorageKey: storageKey,
		StoragePolicyID: session.PolicyID, MimeType: utils.GetMimeType(ext), Status: 1,
	}
	if err := s.fileRepo.Create(newFile); err != nil {
		return nil, fmt.Errorf("保存文件记录失败: %w", err)
	}

	config.DB.Model(session).Updates(map[string]interface{}{"status": 2, "file_id": newFile.ID, "md5": fileMD5})
	return fileToVO(newFile), nil
}

// GetUploadProgress 查询上传进度
func (s *ChunkUploadService) GetUploadProgress(uploadID string) (*dto.ChunkUploadProgressVO, error) {
	var session model.UploadSession
	if err := config.DB.First(&session, uploadID).Error; err != nil {
		return nil, errors.New("上传会话不存在")
	}
	uploaded := parseChunkIndices(session.UploadedChunks)
	missing := []int{}
	for i := 0; i < session.TotalChunks; i++ {
		if !uploaded[i] {
			missing = append(missing, i)
		}
	}
	return &dto.ChunkUploadProgressVO{
		UploadID: session.ID, FileName: session.FileName, FileSize: session.FileSize,
		ChunkSize: session.ChunkSize, TotalChunks: session.TotalChunks,
		Uploaded: len(uploaded), Missing: missing, Status: session.Status,
	}, nil
}

// CancelUpload 取消上传
func (s *ChunkUploadService) CancelUpload(uploadID string) error {
	var session model.UploadSession
	if err := config.DB.First(&session, uploadID).Error; err != nil {
		return errors.New("上传会话不存在")
	}
	os.RemoveAll(filepath.Join(config.AppConfig.Storage.LocalBasePath, ".tmp", uploadID))
	config.DB.Model(&session).Update("status", 3)
	return nil
}

func parseChunkIndices(s string) map[int]bool {
	result := make(map[int]bool)
	if s == "" {
		return result
	}
	for _, part := range strings.Split(s, ",") {
		if idx, err := strconv.Atoi(strings.TrimSpace(part)); err == nil {
			result[idx] = true
		}
	}
	return result
}

func buildChunkIndices(m map[int]bool) string {
	indices := []int{}
	for idx := range m {
		indices = append(indices, idx)
	}
	sort.Ints(indices)
	strs := make([]string, len(indices))
	for i, idx := range indices {
		strs[i] = strconv.Itoa(idx)
	}
	return strings.Join(strs, ",")
}

func fileToVO(file *model.File) *dto.FileVO {
	vo := &dto.FileVO{
		ID: dto.StringID(file.ID), Name: file.Name, Extension: file.Extension, FolderID: dto.StringID(file.FolderID),
		OwnerID: dto.StringID(file.OwnerID),
		Size: file.Size, StorageKey: file.StorageKey, StoragePolicyID: file.StoragePolicyID,
		MimeType: file.MimeType, IsFavorite: file.IsFavorite, DownloadCount: file.DownloadCount,
		Version: file.Version, CreatedAt: file.CreatedAt, UpdatedAt: file.UpdatedAt,
	}
	if file.StorageKey != "" {
		engine, err := storage.GetEngineByPolicyID(file.StoragePolicyID)
		if err == nil {
			if localEngine, ok := engine.(*storage.LocalEngine); ok {
				vo.DownloadURL = localEngine.GetURL(file.StorageKey)
				vo.PreviewURL = localEngine.GetURL(file.StorageKey)
			} else if minioEngine, ok := engine.(*storage.MinIOEngine); ok {
				vo.PreviewURL = minioEngine.GetPreviewURL(file.StorageKey)
			}
		}
	}
	return vo
}