FileRelay/internal/service/batch_service.go

package service

import (
	"FileRelay/internal/bootstrap"
	"FileRelay/internal/model"
	"FileRelay/internal/storage"
	"context"
	"errors"
	"math/big"
	"strings"
	"time"

	"crypto/rand"

	"gorm.io/gorm"
)

type BatchService struct {
	db *gorm.DB
}

func NewBatchService() *BatchService {
	return &BatchService{db: bootstrap.DB}
}

func (s *BatchService) GetBatchByPickupCode(code string) (*model.FileBatch, error) {
	var batch model.FileBatch
	err := s.db.Preload("FileItems").Where("pickup_code = ? AND status = ?", code, "active").First(&batch).Error
	if err != nil {
		return nil, err
	}

	// 检查是否过期
	if s.IsExpired(&batch) {
		s.MarkAsExpired(&batch)
		return nil, errors.New("batch expired")
	}

	return &batch, nil
}

func (s *BatchService) IsExpired(batch *model.FileBatch) bool {
	if batch.Status != "active" {
		return true
	}

	switch batch.ExpireType {
	case "time":
		if batch.ExpireAt != nil && time.Now().After(*batch.ExpireAt) {
			return true
		}
	case "download":
		if batch.MaxDownloads > 0 && batch.DownloadCount >= batch.MaxDownloads {
			return true
		}
	}
	return false
}

func (s *BatchService) MarkAsExpired(batch *model.FileBatch) error {
	return s.db.Model(batch).Update("status", "expired").Error
}

func (s *BatchService) DeleteBatch(ctx context.Context, batchID string) error {
	var batch model.FileBatch
	if err := s.db.Preload("FileItems").First(&batch, "id = ?", batchID).Error; err != nil {
		return err
	}

	// 删除物理文件
	for _, item := range batch.FileItems {
		_ = storage.GlobalStorage.Delete(ctx, item.StoragePath)
	}

	// 删除数据库记录 (软删除 Batch)
	return s.db.Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("batch_id = ?", batch.ID).Delete(&model.FileItem{}).Error; err != nil {
			return err
		}
		if err := tx.Delete(&batch).Error; err != nil {
			return err
		}
		return nil
	})
}

func (s *BatchService) IncrementDownloadCount(batchID string) error {
	return s.db.Model(&model.FileBatch{}).Where("id = ?", batchID).
		UpdateColumn("download_count", gorm.Expr("download_count + ?", 1)).Error
}

func (s *BatchService) 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 = ?", string(b)).Count(&count)
	if count > 0 {
		return s.GeneratePickupCode(length) // 递归生成
	}
	return string(b), nil
}

func (s *BatchService) UpdateAllPickupCodes(newLength int) error {
	var batches []model.FileBatch
	// 只更新未删除的记录，包括 active 和 expired
	if err := s.db.Find(&batches).Error; err != nil {
		return err
	}

	return s.db.Transaction(func(tx *gorm.DB) error {
		for _, batch := range batches {
			oldCode := batch.PickupCode
			if len(oldCode) == newLength {
				continue
			}

			var newCode string
			if len(oldCode) < newLength {
				// 右侧补零，方便用户输入原码后通过补 0 完成输入
				newCode = oldCode + strings.Repeat("0", newLength-len(oldCode))
			} else {
				// 截取前 newLength 位，保留原码头部
				newCode = oldCode[:newLength]
			}

			// 检查冲突 (在事务中检查)
			var count int64
			tx.Model(&model.FileBatch{}).Where("pickup_code = ? AND id != ?", newCode, batch.ID).Count(&count)
			if count > 0 {
				// 如果冲突，生成一个新的随机码
				var err error
				newCode, err = s.generateUniquePickupCodeInTx(tx, newLength)
				if err != nil {
					return err
				}
			}

			if err := tx.Model(&batch).Update("pickup_code", newCode).Error; err != nil {
				return err
			}
		}
		return nil
	})
}

func (s *BatchService) generateUniquePickupCodeInTx(tx *gorm.DB, length int) (string, error) {
	const charset = "0123456789"
	for {
		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
		tx.Model(&model.FileBatch{}).Where("pickup_code = ?", string(b)).Count(&count)
		if count == 0 {
			return string(b), nil
		}
	}
}