coai/auth/redeem.go

package auth

import (
	"chat/globals"
	"chat/utils"
	"database/sql"
	"errors"
	"fmt"
	"github.com/go-redis/redis/v8"
)

type Redeem struct {
	Id    int64   `json:"id"`
	Code  string  `json:"code"`
	Quota float32 `json:"quota"`
	Used  bool    `json:"used"`
}

func GenerateRedeemCodes(db *sql.DB, num int, quota float32) ([]string, error) {
	arr := make([]string, 0)
	idx := 0
	for idx < num {
		code := fmt.Sprintf("nio-%s", utils.GenerateChar(32))
		if err := CreateRedeemCode(db, code, quota); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				continue
			}
			return nil, fmt.Errorf("failed to generate code: %w", err)
		}
		arr = append(arr, code)
		idx++
	}

	return arr, nil
}

func CreateRedeemCode(db *sql.DB, code string, quota float32) error {
	_, err := globals.ExecDb(db, `
		INSERT INTO redeem (code, quota) VALUES (?, ?)
	`, code, quota)
	return err
}

func GetRedeemCode(db *sql.DB, code string) (*Redeem, error) {
	row := globals.QueryRowDb(db, `
		SELECT id, code, quota, used
		FROM redeem
		WHERE code = ?
	`, code)
	var redeem Redeem
	err := row.Scan(&redeem.Id, &redeem.Code, &redeem.Quota, &redeem.Used)

	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, fmt.Errorf("redeem code not found")
		}
		return nil, fmt.Errorf("failed to get redeem code: %w", err)
	}
	return &redeem, nil
}
func (r *Redeem) IsUsed() bool {
	return r.Used
}

func (r *Redeem) Use(db *sql.DB) error {
	_, err := globals.ExecDb(db, `
		UPDATE redeem SET used = TRUE WHERE id = ? AND used = FALSE
	`, r.Id)
	return err
}

func (r *Redeem) GetQuota() float32 {
	return r.Quota
}

func (r *Redeem) UseRedeem(db *sql.DB, user *User) error {
	if r.IsUsed() {
		return fmt.Errorf("this redeem code has been used")
	}

	if err := r.Use(db); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return fmt.Errorf("redeem code not found")
		} else if errors.Is(err, sql.ErrTxDone) {
			return fmt.Errorf("transaction has been closed")
		}
		return fmt.Errorf("failed to use redeem code: %w", err)
	}

	if !user.IncreaseQuota(db, r.GetQuota()) {
		return fmt.Errorf("failed to increase quota for user")
	}

	return nil
}

func (u *User) UseRedeem(db *sql.DB, cache *redis.Client, code string) (float32, error) {
	if useDeeptrain() {
		return 0, errors.New("redeem code is not available in deeptrain mode")
	}

	if redeem, err := GetRedeemCode(db, code); err != nil {
		return 0, err
	} else {
		if err := redeem.UseRedeem(db, u); err != nil {
			return 0, fmt.Errorf("failed to use redeem code: %w", err)
		}

		incrBillingRequest(cache, int64(redeem.GetQuota()*10))
		return redeem.GetQuota(), nil
	}
}