verify/backend/internal/middleware/usage_limit.go

package middleware

import (
	"fmt"
	"time"
	"verification-platform-backend/internal/database"
	"verification-platform-backend/internal/model"
	"verification-platform-backend/pkg/response"

	"github.com/gin-gonic/gin"
)

func CheckApiLimit() gin.HandlerFunc {
	return func(c *gin.Context) {
		userID, exists := c.Get("user_id")
		if !exists {
			response.Error(c, 401, "未授权")
			c.Abort()
			return
		}

		var user model.User
		if err := database.DB.Preload("CurrentPackage").First(&user, userID).Error; err != nil {
			response.Error(c, 500, "获取用户信息失败")
			c.Abort()
			return
		}

		if user.CurrentPackageID == nil {
			response.Error(c, 403, "您还没有购买套餐，请先购买套餐")
			c.Abort()
			return
		}

		var permission model.PackagePermission
		if err := database.DB.Where("package_id = ?", user.CurrentPackageID).First(&permission).Error; err != nil {
			response.Error(c, 500, "获取套餐权限失败")
			c.Abort()
			return
		}

		now := time.Now()
		if user.ApiCallsResetAt == nil || now.Sub(*user.ApiCallsResetAt) >= 24*time.Hour {
			user.ApiCallsUsed = 0
			user.ApiCallsResetAt = &now
			database.DB.Save(&user)
		}

		if user.ApiCallsUsed >= permission.MaxApiCalls {
			response.Error(c, 403, fmt.Sprintf("API调用次数已达上限（%d次/天），请升级套餐", permission.MaxApiCalls))
			c.Abort()
			return
		}

		c.Next()

		user.ApiCallsUsed++
		database.DB.Save(&user)
	}
}

func CheckStorageLimit(fileSize int64) gin.HandlerFunc {
	return func(c *gin.Context) {
		userID, exists := c.Get("user_id")
		if !exists {
			response.Error(c, 401, "未授权")
			c.Abort()
			return
		}

		var user model.User
		if err := database.DB.Preload("CurrentPackage").First(&user, userID).Error; err != nil {
			response.Error(c, 500, "获取用户信息失败")
			c.Abort()
			return
		}

		if user.CurrentPackageID == nil {
			response.Error(c, 403, "您还没有购买套餐，请先购买套餐")
			c.Abort()
			return
		}

		var permission model.PackagePermission
		if err := database.DB.Where("package_id = ?", user.CurrentPackageID).First(&permission).Error; err != nil {
			response.Error(c, 500, "获取套餐权限失败")
			c.Abort()
			return
		}

		maxStorageBytes := int64(permission.MaxStorage) * 1024 * 1024
		if user.StorageUsed+fileSize > maxStorageBytes {
			usedMB := float64(user.StorageUsed) / 1024 / 1024
			maxMB := float64(permission.MaxStorage)
			response.Error(c, 403, fmt.Sprintf("存储空间不足，已使用 %.2f MB / %.2f MB", usedMB, maxMB))
			c.Abort()
			return
		}

		c.Next()
	}
}

func RecordApiUsage() gin.HandlerFunc {
	return func(c *gin.Context) {
		startTime := time.Now()

		c.Next()

		userID, exists := c.Get("user_id")
		if !exists {
			return
		}

		appID, _ := c.Get("app_id")

		duration := time.Since(startTime)

		usage := model.ApiUsage{
			UserID:        userID.(uint),
			ApplicationID: appID.(uint),
			Endpoint:      c.Request.URL.Path,
			Method:        c.Request.Method,
			IPAddress:     c.ClientIP(),
			UserAgent:     c.Request.UserAgent(),
			ResponseTime:  int(duration.Milliseconds()),
			StatusCode:    c.Writer.Status(),
			Success:       c.Writer.Status() < 400,
			CreatedAt:     time.Now(),
		}

		if len(c.Errors) > 0 {
			usage.ErrorMessage = c.Errors.String()
		}

		database.DB.Create(&usage)
	}
}

func UpdateStorageUsed(userID uint, fileSize int64, action string) error {
	var user model.User
	if err := database.DB.First(&user, userID).Error; err != nil {
		return err
	}

	if action == "upload" {
		user.StorageUsed += fileSize
	} else if action == "delete" {
		user.StorageUsed -= fileSize
		if user.StorageUsed < 0 {
			user.StorageUsed = 0
		}
	}

	return database.DB.Save(&user).Error
}