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
}

// AIUserQuotaFindList 分页查询所有用户额度（admin用）
func AIUserQuotaFindList(ctx context.Context, db *gorm.DB, page, pageSize int64) ([]AIUserQuota, int64, error) {
	var list []AIUserQuota
	var total int64

	query := db.WithContext(ctx).Model(&AIUserQuota{})
	if err := query.Count(&total).Error; err != nil {
		return nil, 0, err
	}

	offset := (page - 1) * pageSize
	if offset < 0 {
		offset = 0
	}
	err := query.Order("user_id ASC").Offset(int(offset)).Limit(int(pageSize)).Find(&list).Error
	return list, total, err
}

// AIUserQuotaRecharge 充值（原子操作：balance += amount, total_recharged += amount）
func AIUserQuotaRecharge(ctx context.Context, db *gorm.DB, userId int64, amount float64) error {
	// Ensure record exists
	_, err := AIUserQuotaEnsure(ctx, db, userId)
	if err != nil {
		return err
	}
	return db.WithContext(ctx).Model(&AIUserQuota{}).
		Where("user_id = ?", userId).
		Updates(map[string]interface{}{
			"balance":        gorm.Expr("balance + ?", amount),
			"total_recharged": gorm.Expr("total_recharged + ?", amount),
		}).Error
}

// 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
}