new-api/model/usedata.go.bak-hk88-mergedfix-20260325-093955

package model

import (
	"errors"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// QuotaData 柱状图数据
type QuotaData struct {
	Id        int    `json:"id"`
	UserID    int    `json:"user_id" gorm:"index"`
	Username  string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"`
	ModelName string `json:"model_name" gorm:"index:idx_qdt_model_user_name,priority:1;size:64;default:''"`
	CreatedAt int64  `json:"created_at" gorm:"bigint;index:idx_qdt_created_at,priority:2"`
	TokenUsed int    `json:"token_used" gorm:"default:0"`
	Count     int    `json:"count" gorm:"default:0"`
	Quota     int    `json:"quota" gorm:"default:0"`
}

func UpdateQuotaData() {
	for {
		if common.DataExportEnabled {
			common.SysLog("正在更新数据看板数据...")
			SaveQuotaDataCache()
		}
		sleepMinutes := common.DataExportInterval
		if sleepMinutes <= 0 {
			common.SysLog(fmt.Sprintf("DataExportInterval=%d 非法，回退为 5 分钟", sleepMinutes))
			sleepMinutes = 5
		}
		time.Sleep(time.Duration(sleepMinutes) * time.Minute)
	}
}

var CacheQuotaData = make(map[string]*QuotaData)
var CacheQuotaDataLock = sync.Mutex{}

func quotaDataKey(userId int, username string, modelName string, createdAt int64) string {
	return fmt.Sprintf("%d-%s-%s-%d", userId, username, modelName, createdAt)
}

func logQuotaDataCache(userId int, username string, modelName string, count int, quota int, createdAt int64, tokenUsed int) {
	key := quotaDataKey(userId, username, modelName, createdAt)
	quotaData, ok := CacheQuotaData[key]
	if ok {
		quotaData.Count += count
		quotaData.Quota += quota
		quotaData.TokenUsed += tokenUsed
	} else {
		quotaData = &QuotaData{
			UserID:    userId,
			Username:  username,
			ModelName: modelName,
			CreatedAt: createdAt,
			Count:     count,
			Quota:     quota,
			TokenUsed: tokenUsed,
		}
	}
	CacheQuotaData[key] = quotaData
}

func LogQuotaData(userId int, username string, modelName string, quota int, createdAt int64, tokenUsed int) {
	// 只精确到小时
	createdAt = createdAt - (createdAt % 3600)

	CacheQuotaDataLock.Lock()
	defer CacheQuotaDataLock.Unlock()
	logQuotaDataCache(userId, username, modelName, 1, quota, createdAt, tokenUsed)
}

func SaveQuotaDataCache() {
	CacheQuotaDataLock.Lock()
	if len(CacheQuotaData) == 0 {
		CacheQuotaDataLock.Unlock()
		return
	}
	cacheToFlush := CacheQuotaData
	CacheQuotaData = make(map[string]*QuotaData)
	CacheQuotaDataLock.Unlock()

	size := len(cacheToFlush)
	for _, quotaData := range cacheToFlush {
		if err := upsertQuotaData(quotaData); err != nil {
			common.SysLog(fmt.Sprintf("save quota data cache error: %s", err.Error()))
			CacheQuotaDataLock.Lock()
			logQuotaDataCache(quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.Count, quotaData.Quota, quotaData.CreatedAt, quotaData.TokenUsed)
			CacheQuotaDataLock.Unlock()
		}
	}
	common.SysLog(fmt.Sprintf("保存数据看板数据成功，共保存%d条数据", size))
}

func upsertQuotaData(quotaData *QuotaData) error {
	return DB.Clauses(clause.OnConflict{
		Columns: []clause.Column{
			{Name: "user_id"},
			{Name: "username"},
			{Name: "model_name"},
			{Name: "created_at"},
		},
		DoUpdates: clause.Assignments(map[string]interface{}{
			"count":      gorm.Expr("quota_data.count + EXCLUDED.count"),
			"quota":      gorm.Expr("quota_data.quota + EXCLUDED.quota"),
			"token_used": gorm.Expr("quota_data.token_used + EXCLUDED.token_used"),
		}),
	}).Table("quota_data").Create(quotaData).Error
}

func increaseQuotaData(userId int, username string, modelName string, count int, quota int, createdAt int64, tokenUsed int) {
	err := upsertQuotaData(&QuotaData{
		UserID:    userId,
		Username:  username,
		ModelName: modelName,
		CreatedAt: createdAt,
		Count:     count,
		Quota:     quota,
		TokenUsed: tokenUsed,
	})
	if err != nil {
		common.SysLog(fmt.Sprintf("increaseQuotaData error: %s", err))
	}
}

func normalizeQuotaDataResult(quotaDatas []*QuotaData, err error) ([]*QuotaData, error) {
	if err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return []*QuotaData{}, nil
		}
		return nil, err
	}
	if quotaDatas == nil {
		return []*QuotaData{}, nil
	}
	return quotaDatas, nil
}

func mergeQuotaDataWithCache(quotaDatas []*QuotaData, matcher func(*QuotaData) bool) []*QuotaData {
	merged := make(map[string]*QuotaData, len(quotaDatas))
	for _, item := range quotaDatas {
		if item == nil {
			continue
		}
		copied := *item
		merged[quotaDataKey(copied.UserID, copied.Username, copied.ModelName, copied.CreatedAt)] = &copied
	}

	CacheQuotaDataLock.Lock()
	defer CacheQuotaDataLock.Unlock()
	for _, cached := range CacheQuotaData {
		if cached == nil || !matcher(cached) {
			continue
		}
		key := quotaDataKey(cached.UserID, cached.Username, cached.ModelName, cached.CreatedAt)
		if existing, ok := merged[key]; ok {
			existing.Count += cached.Count
			existing.Quota += cached.Quota
			existing.TokenUsed += cached.TokenUsed
		} else {
			copied := *cached
			merged[key] = &copied
		}
	}

	result := make([]*QuotaData, 0, len(merged))
	for _, item := range merged {
		result = append(result, item)
	}
	sort.Slice(result, func(i, j int) bool {
		if result[i].CreatedAt == result[j].CreatedAt {
			if result[i].ModelName == result[j].ModelName {
				return result[i].UserID < result[j].UserID
			}
			return result[i].ModelName < result[j].ModelName
		}
		return result[i].CreatedAt < result[j].CreatedAt
	})
	return result
}

func GetQuotaDataByUsername(username string, startTime int64, endTime int64) (quotaData []*QuotaData, err error) {
	var quotaDatas []*QuotaData
	err = DB.Table("quota_data").Where("username = ? and created_at >= ? and created_at <= ?", username, startTime, endTime).Find(&quotaDatas).Error
	quotaDatas, err = normalizeQuotaDataResult(quotaDatas, err)
	if err != nil {
		return nil, err
	}
	return mergeQuotaDataWithCache(quotaDatas, func(item *QuotaData) bool {
		return item.Username == username && item.CreatedAt >= startTime && item.CreatedAt <= endTime
	}), nil
}

func GetQuotaDataByUserId(userId int, startTime int64, endTime int64) (quotaData []*QuotaData, err error) {
	var quotaDatas []*QuotaData
	err = DB.Table("quota_data").Where("user_id = ? and created_at >= ? and created_at <= ?", userId, startTime, endTime).Find(&quotaDatas).Error
	quotaDatas, err = normalizeQuotaDataResult(quotaDatas, err)
	if err != nil {
		return nil, err
	}
	return mergeQuotaDataWithCache(quotaDatas, func(item *QuotaData) bool {
		return item.UserID == userId && item.CreatedAt >= startTime && item.CreatedAt <= endTime
	}), nil
}

func GetAllQuotaDates(startTime int64, endTime int64, username string) (quotaData []*QuotaData, err error) {
	if username != "" {
		return GetQuotaDataByUsername(username, startTime, endTime)
	}
	var quotaDatas []*QuotaData
	err = DB.Table("quota_data").Select("model_name, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used, created_at").Where("created_at >= ? and created_at <= ?", startTime, endTime).Group("model_name, created_at").Find(&quotaDatas).Error
	quotaDatas, err = normalizeQuotaDataResult(quotaDatas, err)
	if err != nil {
		return nil, err
	}
	return mergeQuotaDataWithCache(quotaDatas, func(item *QuotaData) bool {
		return item.CreatedAt >= startTime && item.CreatedAt <= endTime
	}), nil
}