base/backend/model/ai_user_quota_model.go

package model

import (
	"context"
	"errors"

	"gorm.io/gorm"
)

// AIUserQuotaFindByUser 根据用户ID查询额度
func AIUserQuotaFindByUser(ctx context.Context, db *gorm.DB, userId int64) (*AIUserQuota, error) {
	var quota AIUserQuota
	result := db.WithContext(ctx).Where("user_id = ?", userId).First(&quota)
	if result.Error != nil {
		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
			return nil, ErrNotFound
		}
		return nil, result.Error
	}
	return &quota, nil
}

// AIUserQuotaEnsure 查找或创建用户额度记录
func AIUserQuotaEnsure(ctx context.Context, db *gorm.DB, userId int64) (*AIUserQuota, error) {
	var quota AIUserQuota
	result := db.WithContext(ctx).Where("user_id = ?", userId).First(&quota)
	if result.Error != nil {
		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
			quota = AIUserQuota{UserId: userId}
			if err := db.WithContext(ctx).Create(&quota).Error; err != nil {
				return nil, err
			}
			return &quota, nil
		}
		return nil, result.Error
	}
	return &quota, nil
}

// AIUserQuotaFreeze 冻结额度（原子操作：balance -= amount, frozen_amount += amount）
func AIUserQuotaFreeze(ctx context.Context, db *gorm.DB, userId int64, amount float64) error {
	result := db.WithContext(ctx).Model(&AIUserQuota{}).
		Where("user_id = ? AND balance >= ?", userId, amount).
		Updates(map[string]interface{}{
			"balance":       gorm.Expr("balance - ?", amount),
			"frozen_amount": gorm.Expr("frozen_amount + ?", amount),
		})
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return errors.New("insufficient balance")
	}
	return nil
}

// AIUserQuotaSettle 结算额度（原子操作：frozen_amount -= frozenAmount, total_consumed += actualCost, balance += refund）
func AIUserQuotaSettle(ctx context.Context, db *gorm.DB, userId int64, frozenAmount, actualCost float64) error {
	refund := frozenAmount - actualCost
	result := db.WithContext(ctx).Model(&AIUserQuota{}).
		Where("user_id = ?", userId).
		Updates(map[string]interface{}{
			"frozen_amount":  gorm.Expr("frozen_amount - ?", frozenAmount),
			"total_consumed": gorm.Expr("total_consumed + ?", actualCost),
			"balance":        gorm.Expr("balance + ?", refund),
		})
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return ErrNotFound
	}
	return nil
}

// AIUserQuotaUnfreeze 解冻额度（原子操作：frozen_amount -= amount, balance += amount）
func AIUserQuotaUnfreeze(ctx context.Context, db *gorm.DB, userId int64, amount float64) error {
	result := db.WithContext(ctx).Model(&AIUserQuota{}).
		Where("user_id = ?", userId).
		Updates(map[string]interface{}{
			"frozen_amount": gorm.Expr("frozen_amount - ?", amount),
			"balance":       gorm.Expr("balance + ?", amount),
		})
	if result.Error != nil {
		return result.Error
	}
	if result.RowsAffected == 0 {
		return ErrNotFound
	}
	return nil
}