doubaochat/backend/services/quota_service.go

package services

import (
	"log"
	"time"

	"d:\项目\dobaochet\backend\config"
	"d:\项目\dobaochet\backend\db"
	"d:\项目\dobaochet\backend\models"
)

// QuotaService 配额管理服务
type QuotaService struct {
	cfg *config.Config
}

// NewQuotaService 创建配额管理服务实例
func NewQuotaService(cfg *config.Config) *QuotaService {
	return &QuotaService{cfg: cfg}
}

// StartQuotaResetJob 启动配额重置定时任务
func (s *QuotaService) StartQuotaResetJob() {
	// 立即执行一次重置检查
	s.ResetExpiredQuotas()

	// 每天检查一次
	ticker := time.NewTicker(24 * time.Hour)
	go func() {
		for range ticker.C {
			s.ResetExpiredQuotas()
		}
	}()
}

// ResetExpiredQuotas 重置过期的配额
func (s *QuotaService) ResetExpiredQuotas() {
	log.Println("开始检查并重置过期配额")

	// 查询所有需要重置配额的用户
	var users []models.User
	result := db.GetDB().Preload("Quota").Where("quota.reset_at <= ?", time.Now()).Find(&users)
	if result.Error != nil {
		log.Printf("查询需要重置配额的用户失败: %v", result.Error)
		return
	}

	if len(users) == 0 {
		log.Println("没有需要重置配额的用户")
		return
	}

	log.Printf("发现 %d 个用户需要重置配额", len(users))

	// 批量重置配额
	tx := db.GetDB().Begin()
	defer func() {
		if r := recover(); r != nil {
			tx.Rollback()
			log.Printf("重置配额时发生 panic: %v", r)
		}
	}()

	for _, user := range users {
		// 重置配额
		user.Quota.UsedTokens = 0
		user.Quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)

		if result := tx.Save(&user.Quota); result.Error != nil {
			tx.Rollback()
			log.Printf("重置用户 %s 配额失败: %v", user.Username, result.Error)
			return
		}

		log.Printf("已重置用户 %s 的配额", user.Username)
	}

	if err := tx.Commit().Error; err != nil {
		log.Printf("提交配额重置事务失败: %v", err)
		return
	}

	log.Printf("成功重置 %d 个用户的配额", len(users))
}

// CheckAndUpdateQuota 检查并更新用户配额
func (s *QuotaService) CheckAndUpdateQuota(userID uint, usedTokens int64) (bool, error) {
	// 查询用户配额
	var quota models.Quota
	result := db.GetDB().Where("user_id = ?", userID).First(&quota)
	if result.Error != nil {
		return false, result.Error
	}

	// 检查配额是否过期
	if quota.ResetAt.Before(time.Now()) {
		// 重置配额
		quota.UsedTokens = 0
		quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)
	}

	// 检查配额是否足够
	if quota.UsedTokens+usedTokens > quota.TotalTokens {
		return false, nil
	}

	// 更新配额
	quota.UsedTokens += usedTokens
	if result := db.GetDB().Save(&quota); result.Error != nil {
		return false, result.Error
	}

	return true, nil
}

// GetUserQuota 获取用户配额信息
func (s *QuotaService) GetUserQuota(userID uint) (*models.Quota, error) {
	var quota models.Quota
	result := db.GetDB().Where("user_id = ?", userID).First(&quota)
	if result.Error != nil {
		return nil, result.Error
	}

	// 检查并重置过期配额
	if quota.ResetAt.Before(time.Now()) {
		quota.UsedTokens = 0
		quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)
		db.GetDB().Save(&quota)
	}

	return &quota, nil
}

// UpdateUserQuota 更新用户配额
func (s *QuotaService) UpdateUserQuota(userID uint, totalTokens, tokenLimit int64) error {
	var quota models.Quota
	result := db.GetDB().Where("user_id = ?", userID).First(&quota)
	if result.Error != nil {
		return result.Error
	}

	// 更新配额
	quota.TotalTokens = totalTokens
	quota.TokenLimit = tokenLimit

	return db.GetDB().Save(&quota).Error
}

// ResetUserQuota 重置单个用户配额
func (s *QuotaService) ResetUserQuota(userID uint) error {
	var quota models.Quota
	result := db.GetDB().Where("user_id = ?", userID).First(&quota)
	if result.Error != nil {
		return result.Error
	}

	// 重置配额
	quota.UsedTokens = 0
	quota.ResetAt = time.Now().Add(s.cfg.Quota.ResetInterval)

	return db.GetDB().Save(&quota).Error
}