heicode-mananger/heicode/model/token.go

package model

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/setting"
	"github.com/heicode/manager/setting/operation_setting"
	"github.com/bytedance/gopkg/util/gopool"
	"gorm.io/gorm"
)

type Token struct {
	Id                 int            `json:"id"`
	UserId             int            `json:"user_id" gorm:"index"`
	Key                string         `json:"key" gorm:"type:varchar(128);uniqueIndex"`
	// KeyHash is the SHA-256 hex digest of the raw `Key` value. H2
	// phase A: dual-write — every Insert and key-update path writes
	// both Key and KeyHash. ValidateUserToken still looks up by Key
	// for now; once telemetry shows zero-NULL coverage we can flip
	// the lookup to KeyHash, drop the Key uniqueIndex, and stop
	// storing plaintext bearers in the DB. This is the upgrade path
	// the V2 device-binding work needs to complete the "no long-lived
	// plaintext credential at rest" promise from the product docs §13.9.
	//
	// Indexed for forward compatibility — phase B switches lookups to
	// this column. Nullable / empty during the rollout window:
	//   - new rows: filled by Token.Insert / UpdateKey paths
	//   - legacy rows: filled by BackfillTokenKeyHash on startup
	// SHA-256 hex == 64 chars; varchar(64) is the natural size.
	KeyHash            string         `json:"-" gorm:"type:varchar(64);index;column:key_hash;default:''"`
	Status             int            `json:"status" gorm:"default:1"`
	Name               string         `json:"name" gorm:"index" `
	CreatedTime        int64          `json:"created_time" gorm:"bigint"`
	AccessedTime       int64          `json:"accessed_time" gorm:"bigint"`
	ExpiredTime        int64          `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
	RemainQuota        int            `json:"remain_quota" gorm:"default:0"`
	UnlimitedQuota     bool           `json:"unlimited_quota"`
	ModelLimitsEnabled bool           `json:"model_limits_enabled"`
	ModelLimits        string         `json:"model_limits" gorm:"type:text"`
	AllowIps           *string        `json:"allow_ips" gorm:"default:''"`
	UsedQuota          int            `json:"used_quota" gorm:"default:0"` // used quota
	Group              string         `json:"group" gorm:"default:''"`
	CrossGroupRetry    bool           `json:"cross_group_retry"` // 跨分组重试，仅auto分组有效
	HideFromUserUI     bool           `json:"hide_from_user_ui" gorm:"default:false"` // 系统托管令牌：列表/API key 详情对用户隐藏（中继仍可用）

	// Device-binding fields (P0 of device-signature work). When DevicePubkey
	// is set, this token is bound to a specific Heicode client install and
	// MUST be presented with X-Heicode-{Device-Id,Timestamp,Nonce,Signature}
	// headers; bare bearer use of the key is rejected. When nil, the token
	// is a legacy bare-bearer token (CLI/SDK back-compat path).
	DeviceId             *string `json:"device_id" gorm:"type:varchar(64);index"`               // client-generated UUID v4
	DevicePubkey         *string `json:"device_pubkey" gorm:"type:text"`                        // base64 32-byte Ed25519 pubkey
	DeviceFingerprint    string  `json:"device_fingerprint" gorm:"type:varchar(64);default:''"` // sha256 of HWID + hostname
	DeviceName           string  `json:"device_name" gorm:"type:varchar(64);default:''"`        // user-friendly: "Chen's MacBook"
	DevicePlatform       string  `json:"device_platform" gorm:"type:varchar(16);default:''"`    // darwin / windows / linux
	DeviceAppVersion     string  `json:"device_app_version" gorm:"type:varchar(32);default:''"` // cc-haha version at pair time
	DeviceBoundAt        int64   `json:"device_bound_at" gorm:"bigint;default:0"`
	DeviceLastSeenIp     string  `json:"device_last_seen_ip" gorm:"type:varchar(45);default:''"`
	DeviceLastUsedAt     int64   `json:"device_last_used_at" gorm:"bigint;default:0"`
	RequireDeviceBinding bool    `json:"require_device_binding" gorm:"default:false"` // per-token override: force signed only
	RevokedAt            int64   `json:"revoked_at" gorm:"bigint;default:0"`
	RevokedReason        string  `json:"revoked_reason" gorm:"type:varchar(128);default:''"`

	DeletedAt gorm.DeletedAt `gorm:"index"`
}

func (token *Token) Clean() {
	token.Key = ""
}

func MaskTokenKey(key string) string {
	if key == "" {
		return ""
	}
	if len(key) <= 4 {
		return strings.Repeat("*", len(key))
	}
	if len(key) <= 8 {
		return key[:2] + "****" + key[len(key)-2:]
	}
	return key[:4] + "**********" + key[len(key)-4:]
}

func (token *Token) GetFullKey() string {
	return token.Key
}

func (token *Token) GetMaskedKey() string {
	return MaskTokenKey(token.Key)
}

func (token *Token) GetIpLimits() []string {
	// delete empty spaces
	//split with \n
	ipLimits := make([]string, 0)
	if token.AllowIps == nil {
		return ipLimits
	}
	cleanIps := strings.ReplaceAll(*token.AllowIps, " ", "")
	if cleanIps == "" {
		return ipLimits
	}
	ips := strings.Split(cleanIps, "\n")
	for _, ip := range ips {
		ip = strings.TrimSpace(ip)
		ip = strings.ReplaceAll(ip, ",", "")
		if ip != "" {
			ipLimits = append(ipLimits, ip)
		}
	}
	return ipLimits
}

func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
	var tokens []*Token
	var err error
	err = DB.Where("user_id = ?", userId).Order("id desc").Limit(num).Offset(startIdx).Find(&tokens).Error
	return tokens, err
}

// GetAllUserTokensVisibleInUI 返回控制台「令牌」列表中应对用户展示的条目（排除系统托管、对用户隐藏的 key）。
func GetAllUserTokensVisibleInUI(userId int, startIdx int, num int) ([]*Token, error) {
	var tokens []*Token
	err := DB.Where("user_id = ? AND hide_from_user_ui = ?", userId, false).
		Order("id desc").Limit(num).Offset(startIdx).Find(&tokens).Error
	return tokens, err
}

// CountUserTokensVisibleInUI 与列表分页一致，仅统计对用户可见的令牌。
func CountUserTokensVisibleInUI(userId int) (int64, error) {
	var total int64
	err := DB.Model(&Token{}).
		Where("user_id = ? AND hide_from_user_ui = ?", userId, false).
		Count(&total).Error
	return total, err
}

// FindTokenByDeviceId is the V2 token-lookup path used when the client
// is authenticated by Ed25519 signature instead of bare bearer. The
// device_id header replaces the role bare sk- played in legacy
// TokenAuth: it's a stable handle the Manager already issued at pair
// time, and the device_pubkey on the row is what verifies the request
// signature.
//
// Returns gorm.ErrRecordNotFound when the device_id doesn't match any
// row — caller should treat as 401.
func FindTokenByDeviceId(deviceId string) (*Token, error) {
	if deviceId == "" {
		return nil, gorm.ErrRecordNotFound
	}
	var t Token
	err := DB.Where("device_id = ? AND device_pubkey IS NOT NULL AND device_pubkey <> ''", deviceId).
		First(&t).Error
	if err != nil {
		return nil, err
	}
	return &t, nil
}

// GetUserDeviceBoundTokens returns the user's device-bound tokens for the
// "Devices" management page. Excludes:
//   - legacy bare-bearer tokens (device_pubkey null OR empty string —
//     both shapes show up from partial-migration rows / direct SQL)
//   - revoked tokens (revoked_at > 0). Users complained that revoked
//     rows lingered in the list as noise ("已撤销, 为什么还有记录"); we
//     keep the soft-delete row in DB for audit + so re-pair can heal
//     it (controller/device.go::PairDevice reactivate branch), but it
//     no longer shows up in the user-facing list.
//
// Sort order: "freshness". A device that has been used recently sorts
// above one that hasn't been used in days. A device that has *never*
// been used (last_used_at = 0) falls back to its pair time, so the
// just-paired device the user expects to see at the top doesn't sink
// to the bottom of the list behind an old-but-actively-used machine.
//
// CASE WHEN is portable across SQLite / MySQL / PostgreSQL (per
// heicode/CLAUDE.md Rule 2). `COALESCE(NULLIF(...))` would also work
// but reads less obviously when the value is an int64 zero rather
// than NULL.
func GetUserDeviceBoundTokens(userId int) ([]*Token, error) {
	var tokens []*Token
	err := DB.Where(
		"user_id = ? AND device_pubkey IS NOT NULL AND device_pubkey <> '' AND revoked_at = 0",
		userId,
	).
		Order("CASE WHEN device_last_used_at = 0 THEN device_bound_at ELSE device_last_used_at END DESC, id DESC").
		Find(&tokens).Error
	return tokens, err
}

// CountUserDeviceBoundTokens counts ACTIVE (non-revoked) device-bound
// tokens for the per-user cap enforcement at pair time.
func CountUserDeviceBoundTokens(userId int) (int64, error) {
	var total int64
	err := DB.Model(&Token{}).
		Where("user_id = ? AND device_pubkey IS NOT NULL AND device_pubkey <> '' AND revoked_at = 0", userId).
		Count(&total).Error
	return total, err
}

// UpdateTokenDeviceLastSeen records the most-recent IP and timestamp this
// device-bound token was authenticated with. Called from a goroutine after
// successful signature verification — should be best-effort, never block.
// Nil-guards DB so unit tests (or partial-init binaries) don't panic.
func UpdateTokenDeviceLastSeen(tokenId int, ip string, nowMs int64) error {
	if DB == nil {
		return nil
	}
	return DB.Model(&Token{}).Where("id = ?", tokenId).Updates(map[string]interface{}{
		"device_last_seen_ip": ip,
		"device_last_used_at": nowMs,
	}).Error
}

// RevokeDevice soft-disables a device-bound token. We don't hard-delete
// so the audit trail (last seen IP, fingerprint) stays inspectable.
func RevokeDevice(tokenId int, reason string) error {
	now := common.GetTimestamp() * 1000
	return DB.Model(&Token{}).Where("id = ?", tokenId).Updates(map[string]interface{}{
		"status":         common.TokenStatusDisabled,
		"revoked_at":     now,
		"revoked_reason": reason,
	}).Error
}

// ReactivateDevice undoes RevokeDevice for the case where the SAME device
// (same device_id + same pubkey) re-pairs after being revoked. The pair
// controller has already verified user ownership and pubkey match; this
// only clears the soft-disable state. We deliberately keep the original
// device_fingerprint / device_bound_at — the device IS the same one, and
// rewriting bound_at would corrupt the audit timeline.
//
// `gorm.Updates` with a map skips zero values, but we want to explicitly
// set revoked_at back to 0 and revoked_reason back to '' — so we use
// `Select` to force them through. Same trick the rest of the codebase
// uses when zeroing GORM fields (e.g. ResetUserPassword in user.go).
func ReactivateDevice(tokenId int) error {
	if DB == nil {
		return nil
	}
	return DB.Model(&Token{}).Where("id = ?", tokenId).
		Select("status", "revoked_at", "revoked_reason").
		Updates(map[string]interface{}{
			"status":         common.TokenStatusEnabled,
			"revoked_at":     int64(0),
			"revoked_reason": "",
		}).Error
}

// EnsureUserRelayToken 在用户没有任何 relay 令牌时创建一枚系统托管令牌（对用户隐藏），用于登录后即可走网关。
func EnsureUserRelayToken(userId int, username string) {
	if userId <= 0 {
		return
	}
	count, err := CountUserTokens(userId)
	if err != nil {
		common.SysLog("EnsureUserRelayToken count: " + err.Error())
		return
	}
	if count > 0 {
		return
	}
	key, err := common.GenerateKey()
	if err != nil {
		common.SysLog("EnsureUserRelayToken generate key: " + err.Error())
		return
	}
	token := Token{
		UserId:             userId,
		Name:               username + "的初始令牌",
		Key:                key,
		CreatedTime:        common.GetTimestamp(),
		AccessedTime:       common.GetTimestamp(),
		ExpiredTime:        -1,
		RemainQuota:        500000,
		UnlimitedQuota:     true,
		ModelLimitsEnabled: false,
		HideFromUserUI:     true,
	}
	if setting.DefaultUseAutoGroup {
		token.Group = "auto"
	}
	if err := token.Insert(); err != nil {
		common.SysLog("EnsureUserRelayToken insert: " + err.Error())
	}
}

// sanitizeLikePattern 校验并清洗用户输入的 LIKE 搜索模式。
// 规则：
//  1. 转义 ! 和 _（使用 ! 作为 ESCAPE 字符，兼容 MySQL/PostgreSQL/SQLite）
//  2. 连续的 % 合并为单个 %
//  3. 最多允许 2 个 %
//  4. 含 % 时（模糊搜索），去掉 % 后关键词长度必须 >= 2
//  5. 不含 % 时按精确匹配
func sanitizeLikePattern(input string) (string, error) {
	// 1. 先转义 ESCAPE 字符 ! 自身，再转义 _
	//    使用 ! 而非 \ 作为 ESCAPE 字符，避免 MySQL 中反斜杠的字符串转义问题
	input = strings.ReplaceAll(input, "!", "!!")
	input = strings.ReplaceAll(input, `_`, `!_`)

	// 2. 连续的 % 直接拒绝
	if strings.Contains(input, "%%") {
		return "", errors.New("搜索模式中不允许包含连续的 % 通配符")
	}

	// 3. 统计 % 数量，不得超过 2
	count := strings.Count(input, "%")
	if count > 2 {
		return "", errors.New("搜索模式中最多允许包含 2 个 % 通配符")
	}

	// 4. 含 % 时，去掉 % 后关键词长度必须 >= 2
	if count > 0 {
		stripped := strings.ReplaceAll(input, "%", "")
		if len(stripped) < 2 {
			return "", errors.New("使用模糊搜索时，关键词长度至少为 2 个字符")
		}
		return input, nil
	}

	// 5. 无 % 时，精确全匹配
	return input, nil
}

const searchHardLimit = 100

// SearchUserTokens 搜索用户令牌。restrictVisibleInUI 为 true 时排除系统托管、对用户隐藏的令牌（与普通用户控制台列表一致）。
func SearchUserTokens(userId int, keyword string, token string, offset int, limit int, restrictVisibleInUI bool) (tokens []*Token, total int64, err error) {
	// model 层强制截断
	if limit <= 0 || limit > searchHardLimit {
		limit = searchHardLimit
	}
	if offset < 0 {
		offset = 0
	}

	if token != "" {
		token = strings.TrimPrefix(token, "sk-")
	}

	// 超量用户（令牌数超过上限）只允许精确搜索，禁止模糊搜索
	maxTokens := operation_setting.GetMaxUserTokens()
	hasFuzzy := strings.Contains(keyword, "%") || strings.Contains(token, "%")
	if hasFuzzy {
		var count int64
		var err error
		if restrictVisibleInUI {
			count, err = CountUserTokensVisibleInUI(userId)
		} else {
			count, err = CountUserTokens(userId)
		}
		if err != nil {
			common.SysLog("failed to count user tokens: " + err.Error())
			return nil, 0, errors.New("获取令牌数量失败")
		}
		if int(count) > maxTokens {
			return nil, 0, errors.New("令牌数量超过上限，仅允许精确搜索，请勿使用 % 通配符")
		}
	}

	baseQuery := DB.Model(&Token{}).Where("user_id = ?", userId)
	if restrictVisibleInUI {
		baseQuery = baseQuery.Where("hide_from_user_ui = ?", false)
	}

	// 非空才加 LIKE 条件，空则跳过（不过滤该字段）
	if keyword != "" {
		keywordPattern, err := sanitizeLikePattern(keyword)
		if err != nil {
			return nil, 0, err
		}
		baseQuery = baseQuery.Where("name LIKE ? ESCAPE '!'", keywordPattern)
	}
	if token != "" {
		tokenPattern, err := sanitizeLikePattern(token)
		if err != nil {
			return nil, 0, err
		}
		baseQuery = baseQuery.Where(commonKeyCol+" LIKE ? ESCAPE '!'", tokenPattern)
	}

	// 先查匹配总数（用于分页，受 maxTokens 上限保护，避免全表 COUNT）
	err = baseQuery.Limit(maxTokens).Count(&total).Error
	if err != nil {
		common.SysError("failed to count search tokens: " + err.Error())
		return nil, 0, errors.New("搜索令牌失败")
	}

	// 再分页查数据
	err = baseQuery.Order("id desc").Offset(offset).Limit(limit).Find(&tokens).Error
	if err != nil {
		common.SysError("failed to search tokens: " + err.Error())
		return nil, 0, errors.New("搜索令牌失败")
	}
	return tokens, total, nil
}

func ValidateUserToken(key string) (token *Token, err error) {
	if key == "" {
		return nil, ErrTokenNotProvided
	}
	token, err = GetTokenByKey(key, false)
	if err == nil {
		if token.Status == common.TokenStatusExhausted ||
			token.Status == common.TokenStatusExpired ||
			token.Status != common.TokenStatusEnabled {
			return token, ErrTokenInvalid
		}
		if token.ExpiredTime != -1 && token.ExpiredTime < common.GetTimestamp() {
			if !common.RedisEnabled {
				token.Status = common.TokenStatusExpired
				err := token.SelectUpdate()
				if err != nil {
					common.SysLog("failed to update token status" + err.Error())
				}
			}
			return token, ErrTokenInvalid
		}
		if !token.UnlimitedQuota && token.RemainQuota <= 0 {
			if !common.RedisEnabled {
				token.Status = common.TokenStatusExhausted
				err := token.SelectUpdate()
				if err != nil {
					common.SysLog("failed to update token status" + err.Error())
				}
			}
			return token, ErrTokenInvalid
		}
		return token, nil
	}
	common.SysLog("ValidateUserToken: failed to get token: " + err.Error())
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, ErrTokenInvalid
	}
	return nil, fmt.Errorf("%w: %v", ErrDatabase, err)
}

func GetTokenByIds(id int, userId int) (*Token, error) {
	if id == 0 || userId == 0 {
		return nil, errors.New("id 或 userId 为空！")
	}
	token := Token{Id: id, UserId: userId}
	var err error = nil
	err = DB.First(&token, "id = ? and user_id = ?", id, userId).Error
	return &token, err
}

func GetTokenById(id int) (*Token, error) {
	if id == 0 {
		return nil, errors.New("id 为空！")
	}
	token := Token{Id: id}
	var err error = nil
	err = DB.First(&token, "id = ?", id).Error
	if shouldUpdateRedis(true, err) {
		gopool.Go(func() {
			if err := cacheSetToken(token); err != nil {
				common.SysLog("failed to update user status cache: " + err.Error())
			}
		})
	}
	return &token, err
}

func GetTokenByKey(key string, fromDB bool) (token *Token, err error) {
	defer func() {
		// Update Redis cache asynchronously on successful DB read
		if shouldUpdateRedis(fromDB, err) && token != nil {
			gopool.Go(func() {
				if err := cacheSetToken(*token); err != nil {
					common.SysLog("failed to update user status cache: " + err.Error())
				}
			})
		}
	}()
	if !fromDB && common.RedisEnabled {
		// Try Redis first
		token, err := cacheGetTokenByKey(key)
		if err == nil {
			return token, nil
		}
		// Don't return error - fall through to DB
	}
	fromDB = true
	err = DB.Where(commonKeyCol+" = ?", key).First(&token).Error
	return token, err
}

// BackfillTokenKeyHash fills tokens.key_hash for any rows that still
// have an empty hash but a non-empty plaintext Key. Idempotent — safe
// to call on every container startup. Batches rows so a large legacy
// `tokens` table (10k+ rows) doesn't blow up memory.
//
// This is the H2 phase-A backfill. Once telemetry shows
// COUNT(*) WHERE key_hash = '' AND key != '' equals zero across all
// environments, phase B switches ValidateUserToken to look up by hash
// and we can drop the Key column.
func BackfillTokenKeyHash() error {
	if DB == nil {
		return nil
	}
	const batchSize = 500
	totalUpdated := 0
	for {
		var rows []Token
		// Select only id + key, skip deleted-at rows, only the ones
		// that still need a hash. The model's regular Find honours
		// soft-delete; that's correct — we don't backfill tombstones.
		err := DB.Select("id", "key").
			Where("(key_hash IS NULL OR key_hash = '') AND key IS NOT NULL AND key <> ''").
			Limit(batchSize).
			Find(&rows).Error
		if err != nil {
			return err
		}
		if len(rows) == 0 {
			break
		}
		for _, r := range rows {
			hash := computeKeyHash(r.Key)
			if hash == "" {
				continue
			}
			if err := DB.Model(&Token{}).Where("id = ?", r.Id).
				Update("key_hash", hash).Error; err != nil {
				return err
			}
			totalUpdated++
		}
		if len(rows) < batchSize {
			break
		}
	}
	if totalUpdated > 0 {
		common.SysLog(fmt.Sprintf("BackfillTokenKeyHash: filled %d rows", totalUpdated))
	}
	return nil
}

// computeKeyHash returns the SHA-256 hex digest of `key`. H2 phase A —
// stored in tokens.key_hash alongside the plaintext Key so we can move
// the auth lookup off plaintext in phase B without a downtime
// migration. Empty key returns empty string (zero-value, intentional).
func computeKeyHash(key string) string {
	if key == "" {
		return ""
	}
	sum := sha256.Sum256([]byte(key))
	return hex.EncodeToString(sum[:])
}

func (token *Token) Insert() error {
	// H2: always populate KeyHash on write so phase B can flip the
	// lookup index without a one-off backfill round-trip. Safe to
	// re-set on every Insert — value is deterministic from Key.
	token.KeyHash = computeKeyHash(token.Key)
	return DB.Create(token).Error
}

// Update Make sure your token's fields is completed, because this will update non-zero values
func (token *Token) Update() (err error) {
	defer func() {
		if shouldUpdateRedis(true, err) {
			gopool.Go(func() {
				err := cacheSetToken(*token)
				if err != nil {
					common.SysLog("failed to update token cache: " + err.Error())
				}
			})
		}
	}()
	err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
		"model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error
	return err
}

func (token *Token) SelectUpdate() (err error) {
	defer func() {
		if shouldUpdateRedis(true, err) {
			gopool.Go(func() {
				err := cacheSetToken(*token)
				if err != nil {
					common.SysLog("failed to update token cache: " + err.Error())
				}
			})
		}
	}()
	// This can update zero values
	return DB.Model(token).Select("accessed_time", "status").Updates(token).Error
}

func (token *Token) Delete() (err error) {
	defer func() {
		if shouldUpdateRedis(true, err) {
			gopool.Go(func() {
				err := cacheDeleteToken(token.Key)
				if err != nil {
					common.SysLog("failed to delete token cache: " + err.Error())
				}
			})
		}
	}()
	err = DB.Delete(token).Error
	return err
}

func (token *Token) IsModelLimitsEnabled() bool {
	return token.ModelLimitsEnabled
}

func (token *Token) GetModelLimits() []string {
	if token.ModelLimits == "" {
		return []string{}
	}
	return strings.Split(token.ModelLimits, ",")
}

func (token *Token) GetModelLimitsMap() map[string]bool {
	limits := token.GetModelLimits()
	limitsMap := make(map[string]bool)
	for _, limit := range limits {
		limitsMap[limit] = true
	}
	return limitsMap
}

func DisableModelLimits(tokenId int) error {
	token, err := GetTokenById(tokenId)
	if err != nil {
		return err
	}
	token.ModelLimitsEnabled = false
	token.ModelLimits = ""
	return token.Update()
}

func DeleteTokenById(id int, userId int) (err error) {
	// Why we need userId here? In case user want to delete other's token.
	if id == 0 || userId == 0 {
		return errors.New("id 或 userId 为空！")
	}
	token := Token{Id: id, UserId: userId}
	err = DB.Where(token).First(&token).Error
	if err != nil {
		return err
	}
	return token.Delete()
}

func IncreaseTokenQuota(tokenId int, key string, quota int) (err error) {
	if quota < 0 {
		return errors.New("quota 不能为负数！")
	}
	if common.RedisEnabled {
		gopool.Go(func() {
			err := cacheIncrTokenQuota(key, int64(quota))
			if err != nil {
				common.SysLog("failed to increase token quota: " + err.Error())
			}
		})
	}
	if common.BatchUpdateEnabled {
		addNewRecord(BatchUpdateTypeTokenQuota, tokenId, quota)
		return nil
	}
	return increaseTokenQuota(tokenId, quota)
}

func increaseTokenQuota(id int, quota int) (err error) {
	err = DB.Model(&Token{}).Where("id = ?", id).Updates(
		map[string]interface{}{
			"remain_quota":  gorm.Expr("remain_quota + ?", quota),
			"used_quota":    gorm.Expr("used_quota - ?", quota),
			"accessed_time": common.GetTimestamp(),
		},
	).Error
	return err
}

func DecreaseTokenQuota(id int, key string, quota int) (err error) {
	if quota < 0 {
		return errors.New("quota 不能为负数！")
	}
	if common.RedisEnabled {
		gopool.Go(func() {
			err := cacheDecrTokenQuota(key, int64(quota))
			if err != nil {
				common.SysLog("failed to decrease token quota: " + err.Error())
			}
		})
	}
	if common.BatchUpdateEnabled {
		addNewRecord(BatchUpdateTypeTokenQuota, id, -quota)
		return nil
	}
	return decreaseTokenQuota(id, quota)
}

func decreaseTokenQuota(id int, quota int) (err error) {
	err = DB.Model(&Token{}).Where("id = ?", id).Updates(
		map[string]interface{}{
			"remain_quota":  gorm.Expr("remain_quota - ?", quota),
			"used_quota":    gorm.Expr("used_quota + ?", quota),
			"accessed_time": common.GetTimestamp(),
		},
	).Error
	return err
}

// CountUserTokens returns total number of tokens for the given user, used for pagination
func CountUserTokens(userId int) (int64, error) {
	var total int64
	err := DB.Model(&Token{}).Where("user_id = ?", userId).Count(&total).Error
	return total, err
}

// BatchDeleteTokens 删除指定用户的一组令牌，返回成功删除数量
func BatchDeleteTokens(ids []int, userId int) (int, error) {
	if len(ids) == 0 {
		return 0, errors.New("ids 不能为空！")
	}

	tx := DB.Begin()

	var tokens []Token
	if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Find(&tokens).Error; err != nil {
		tx.Rollback()
		return 0, err
	}

	if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Delete(&Token{}).Error; err != nil {
		tx.Rollback()
		return 0, err
	}

	if err := tx.Commit().Error; err != nil {
		return 0, err
	}

	if common.RedisEnabled {
		gopool.Go(func() {
			for _, t := range tokens {
				_ = cacheDeleteToken(t.Key)
			}
		})
	}

	return len(tokens), nil
}

func GetTokenKeysByIds(ids []int, userId int) ([]Token, error) {
	var tokens []Token
	err := DB.Select("id", commonKeyCol).
		Where("user_id = ? AND id IN (?)", userId, ids).
		Find(&tokens).Error
	return tokens, err
}

// InvalidateUserTokensCache 清理指定用户所有令牌在 Redis 中的缓存，
// 配合 InvalidateUserCache 使用，可在用户被禁用/删除时立即阻断其令牌的请求。
// 下一次请求将从数据库重新加载令牌及用户状态，从而立即识别出被禁用的用户。
func InvalidateUserTokensCache(userId int) error {
	if !common.RedisEnabled {
		return nil
	}
	if userId <= 0 {
		return errors.New("userId 无效")
	}
	var tokens []Token
	if err := DB.Unscoped().
		Select("id", commonKeyCol).
		Where("user_id = ?", userId).
		Find(&tokens).Error; err != nil {
		return err
	}
	var firstErr error
	for _, t := range tokens {
		if t.Key == "" {
			continue
		}
		if err := cacheDeleteToken(t.Key); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	return firstErr
}