FileRelay/internal/service/upload_service.go

package service

import (
	"FileRelay/internal/bootstrap"
	"FileRelay/internal/config"
	"FileRelay/internal/model"
	"FileRelay/internal/storage"
	"context"
	"crypto/rand"
	"fmt"
	"math/big"
	"mime/multipart"
	"path/filepath"
	"time"

	"github.com/google/uuid"
	"gorm.io/gorm"
)

type UploadService struct {
	db *gorm.DB
}

func NewUploadService() *UploadService {
	return &UploadService{db: bootstrap.DB}
}

func (s *UploadService) CreateBatch(ctx context.Context, files []*multipart.FileHeader, remark string, expireType string, expireValue interface{}) (*model.FileBatch, error) {
	// 1. 生成取件码
	pickupCode, err := s.generatePickupCode(config.GlobalConfig.Security.PickupCodeLength)
	if err != nil {
		return nil, err
	}

	// 2. 准备 Batch
	batch := &model.FileBatch{
		ID:         uuid.New().String(),
		PickupCode: pickupCode,
		Remark:     remark,
		ExpireType: expireType,
		Status:     "active",
		Type:       "file",
	}

	s.applyExpire(batch, expireType, expireValue)

	// 3. 处理文件上传
	return batch, s.db.Transaction(func(tx *gorm.DB) error {
		if err := tx.Create(batch).Error; err != nil {
			return err
		}

		for _, fileHeader := range files {
			fileItem, err := s.processFile(ctx, tx, batch.ID, fileHeader)
			if err != nil {
				return err
			}
			batch.FileItems = append(batch.FileItems, *fileItem)
		}
		return nil
	})
}

func (s *UploadService) CreateTextBatch(ctx context.Context, content string, remark string, expireType string, expireValue interface{}) (*model.FileBatch, error) {
	// 1. 生成取件码
	pickupCode, err := s.generatePickupCode(config.GlobalConfig.Security.PickupCodeLength)
	if err != nil {
		return nil, err
	}

	// 2. 准备 Batch
	batch := &model.FileBatch{
		ID:         uuid.New().String(),
		PickupCode: pickupCode,
		Remark:     remark,
		ExpireType: expireType,
		Status:     "active",
		Type:       "text",
		Content:    content,
	}

	s.applyExpire(batch, expireType, expireValue)

	// 3. 保存
	if err := s.db.Create(batch).Error; err != nil {
		return nil, err
	}

	return batch, nil
}

func (s *UploadService) applyExpire(batch *model.FileBatch, expireType string, expireValue interface{}) {
	switch expireType {
	case "time":
		if days, ok := expireValue.(int); ok {
			expireAt := time.Now().Add(time.Duration(days) * 24 * time.Hour)
			batch.ExpireAt = &expireAt
		}
	case "download":
		if max, ok := expireValue.(int); ok {
			batch.MaxDownloads = max
		}
	}
}

func (s *UploadService) processFile(ctx context.Context, tx *gorm.DB, batchID string, fileHeader *multipart.FileHeader) (*model.FileItem, error) {
	file, err := fileHeader.Open()
	if err != nil {
		return nil, err
	}
	defer file.Close()

	// 生成唯一存储路径
	ext := filepath.Ext(fileHeader.Filename)
	fileID := uuid.New().String()
	storagePath := fmt.Sprintf("%s/%s%s", batchID, fileID, ext)

	// 保存到存储层
	if err := storage.GlobalStorage.Save(ctx, storagePath, file); err != nil {
		return nil, err
	}

	// 创建数据库记录
	item := &model.FileItem{
		ID:           fileID,
		BatchID:      batchID,
		OriginalName: fileHeader.Filename,
		StoragePath:  storagePath,
		Size:         fileHeader.Size,
		MimeType:     fileHeader.Header.Get("Content-Type"),
	}

	if err := tx.Create(item).Error; err != nil {
		return nil, err
	}

	return item, nil
}

func (s *UploadService) generatePickupCode(length int) (string, error) {
	const charset = "0123456789"
	b := make([]byte, length)
	for i := range b {
		num, err := rand.Int(rand.Reader, big.NewInt(int64(len(charset))))
		if err != nil {
			return "", err
		}
		b[i] = charset[num.Int64()]
	}
	// 检查是否冲突
	var count int64
	s.db.Model(&model.FileBatch{}).Where("pickup_code = ? AND status = ?", string(b), "active").Count(&count)
	if count > 0 {
		return s.generatePickupCode(length) // 递归生成
	}
	return string(b), nil
}