opencatd-open/backend/internal/dao/usage.go

package dao

import (
	"context"
	"opencatd-open/internal/store"
	"time"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type UsageDAO struct {
	db *gorm.DB
}

type DailyUsageDAO struct {
	db *gorm.DB
}

func NewUsageDAO(db *gorm.DB) *UsageDAO {
	return &UsageDAO{db: db}
}

func NewDailyUsageDAO(db *gorm.DB) *DailyUsageDAO {
	return &DailyUsageDAO{db: db}
}

// UsageLog DAO
func (d *UsageDAO) Create(ctx context.Context, log *store.UsageLog) error {
	return d.db.WithContext(ctx).Create(log).Error
}

func (d *UsageDAO) BatchCreate(ctx context.Context, logs []*store.UsageLog) error {
	return d.db.WithContext(ctx).Create(logs).Error
}

func (d *UsageDAO) ListByUserID(ctx context.Context, userID uint64, limit, offset int) ([]*store.UsageLog, error) {
	var logs []*store.UsageLog
	err := d.db.WithContext(ctx).
		Where("user_id = ?", userID).
		Order("created_at DESC").
		Limit(limit).
		Offset(offset).
		Find(&logs).Error
	return logs, err
}

func (d *UsageDAO) Delete(ctx context.Context, id uint64) error {
	return d.db.WithContext(ctx).Delete(&store.UsageLog{}, id).Error
}

func (d *UsageDAO) CountByUserID(ctx context.Context, userID uint64) (int64, error) {
	var count int64
	err := d.db.WithContext(ctx).Model(&store.UsageLog{}).Where("user_id = ?", userID).Count(&count).Error
	return count, err
}

// UsageFilter 用量明细筛选条件（管理后台）。
type UsageFilter struct {
	UserID    *uint64 // 指定用户（nil=全部）
	Protocol  string  // 协议 chat/messages/responses（空=全部）
	Status    string  // success/error/canceled（空=全部）
	ModelName string  // 模型名模糊（空=全部）
	Limit     int
	Offset    int
}

// ListAll 管理后台全量用量明细（分页 + 筛选），并带用户名。
func (d *UsageDAO) ListAll(ctx context.Context, f UsageFilter) ([]store.UsageLog, error) {
	q := d.db.WithContext(ctx).Model(&store.UsageLog{})
	q = applyUsageFilter(q, f)
	var logs []store.UsageLog
	err := q.Order("created_at DESC").Limit(f.Limit).Offset(f.Offset).Find(&logs).Error
	return logs, err
}

// CountAll 统计符合筛选条件的明细总数。
func (d *UsageDAO) CountAll(ctx context.Context, f UsageFilter) (int64, error) {
	q := d.db.WithContext(ctx).Model(&store.UsageLog{})
	q = applyUsageFilter(q, f)
	var count int64
	err := q.Count(&count).Error
	return count, err
}

func applyUsageFilter(q *gorm.DB, f UsageFilter) *gorm.DB {
	if f.UserID != nil {
		q = q.Where("user_id = ?", *f.UserID)
	}
	if f.Protocol != "" {
		q = q.Where("protocol = ?", f.Protocol)
	}
	if f.Status != "" {
		q = q.Where("status = ?", f.Status)
	}
	if f.ModelName != "" {
		q = q.Where("model_name LIKE ?", "%"+f.ModelName+"%")
	}
	return q
}

// UsageDaily DAO
func (d *DailyUsageDAO) Create(ctx context.Context, log *store.UsageDaily) error {
	return d.db.WithContext(ctx).Create(log).Error
}

func (d *DailyUsageDAO) ListByUserID(ctx context.Context, userID uint64, limit, offset int) ([]*store.UsageDaily, error) {
	var logs []*store.UsageDaily
	err := d.db.WithContext(ctx).
		Where("user_id = ?", userID).
		Order("date DESC").
		Limit(limit).
		Offset(offset).
		Find(&logs).Error
	return logs, err
}

func (d *DailyUsageDAO) GetByDate(ctx context.Context, userID uint64, date string) (*store.UsageDaily, error) {
	var log store.UsageDaily
	err := d.db.WithContext(ctx).
		Where("user_id = ? AND date = ?", userID, date).
		First(&log).Error
	if err != nil {
		return nil, err
	}
	return &log, nil
}

// UpsertDailyUsage 按 (user_id, model_id, date) 累加式 upsert：
// 行不存在则插入；存在则在原值基础上增量累加（不能用 AssignmentColumns 覆盖，
// 否则多次 flush 会互相清零）。非限定列名在 SQLite/MySQL/PG 的 upsert 语义下都指向目标行。
func (d *DailyUsageDAO) UpsertDailyUsage(ctx context.Context, log *store.UsageDaily) error {
	return d.db.WithContext(ctx).Clauses(clause.OnConflict{
		Columns: []clause.Column{{Name: "user_id"}, {Name: "model_id"}, {Name: "date"}},
		DoUpdates: clause.Assignments(map[string]interface{}{
			"requests":          gorm.Expr("requests + ?", log.Requests),
			"input_tokens":      gorm.Expr("input_tokens + ?", log.InputTokens),
			"output_tokens":     gorm.Expr("output_tokens + ?", log.OutputTokens),
			"cache_read_tokens": gorm.Expr("cache_read_tokens + ?", log.CacheReadTokens),
			"cost":              gorm.Expr("cost + ?", log.Cost),
		}),
	}).Create(log).Error
}

func (d *DailyUsageDAO) ListByDateRange(ctx context.Context, userID uint64, start, end time.Time) ([]*store.UsageDaily, error) {
	var logs []*store.UsageDaily
	err := d.db.WithContext(ctx).
		Where("user_id = ? AND date >= ? AND date <= ?", userID, start.Format("2006-01-02"), end.Format("2006-01-02")).
		Order("date DESC").
		Find(&logs).Error
	return logs, err
}

// ListAll 管理后台：全部用户的日聚合（可选按用户/日期范围筛选），按日期倒序。
func (d *DailyUsageDAO) ListAll(ctx context.Context, userID *uint64, start, end string) ([]store.UsageDaily, error) {
	q := d.db.WithContext(ctx).Model(&store.UsageDaily{})
	if userID != nil {
		q = q.Where("user_id = ?", *userID)
	}
	if start != "" {
		q = q.Where("date >= ?", start)
	}
	if end != "" {
		q = q.Where("date <= ?", end)
	}
	var logs []store.UsageDaily
	err := q.Order("date DESC, user_id ASC").Find(&logs).Error
	return logs, err
}