openteam/server/internal/usage/handler.go

package usage

import (
	"net/http"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/shopspring/decimal"
	"go.uber.org/zap"
	"gorm.io/gorm"

	"openteam/server/internal/pkg/httpx"
	"openteam/server/internal/store"
	"openteam/server/internal/user"
)

type Handler struct {
	db  *gorm.DB
	log *zap.Logger
}

func NewHandler(db *gorm.DB, log *zap.Logger) *Handler {
	return &Handler{db: db, log: log}
}

// Summary handles GET /api/v1/usage/summary.
func (h *Handler) Summary(c *gin.Context) {
	u := user.Current(c)
	now := time.Now()
	today := now.Format("2006-01-02")
	month := now.Format("2006-01")

	type agg struct {
		Requests    int
		InputTokens int64
		OutputTokens int64
		Cost        decimal.Decimal
	}

	todayAgg := h.aggregate(u.ID, today, today)
	monthAgg := h.aggregate(u.ID, month+"-01", now.Format("2006-01-02"))
	total := h.aggregate(u.ID, "", "")

	httpx.OK(c, gin.H{
		"today": gin.H{
			"requests":     todayAgg.Requests,
			"inputTokens":  todayAgg.InputTokens,
			"outputTokens": todayAgg.OutputTokens,
			"cost":         todayAgg.Cost.String(),
		},
		"month": gin.H{
			"requests":     monthAgg.Requests,
			"inputTokens":  monthAgg.InputTokens,
			"outputTokens": monthAgg.OutputTokens,
			"cost":         monthAgg.Cost.String(),
		},
		"total": gin.H{
			"requests":     total.Requests,
			"inputTokens":  total.InputTokens,
			"outputTokens": total.OutputTokens,
			"cost":         total.Cost.String(),
		},
	})
}

func (h *Handler) aggregate(userID int64, from, to string) struct {
	Requests    int
	InputTokens int64
	OutputTokens int64
	Cost        decimal.Decimal
} {
	var out struct {
		Requests    int
		InputTokens int64
		OutputTokens int64
		Cost        decimal.Decimal
	}
	q := h.db.Model(&store.UsageDaily{}).Where("user_id = ?", userID)
	if from != "" {
		q = q.Where("date >= ?", from)
	}
	if to != "" {
		q = q.Where("date <= ?", to)
	}
	q.Select("COALESCE(SUM(requests),0) as requests, COALESCE(SUM(input_tokens),0) as input_tokens, COALESCE(SUM(output_tokens),0) as output_tokens, COALESCE(SUM(cost),0) as cost").
		Scan(&out)
	return out
}

// Stats handles GET /api/v1/usage/stats?from&to&group=day|model.
func (h *Handler) Stats(c *gin.Context) {
	u := user.Current(c)
	from := c.Query("from")
	to := c.Query("to")
	group := c.DefaultQuery("group", "day")

	var rows []struct {
		Key          string `gorm:"column:g"`
		Requests     int
		InputTokens  int64
		OutputTokens int64
		Cost         decimal.Decimal
	}

	q := h.db.Model(&store.UsageDaily{}).Where("user_id = ?", u.ID)
	if from != "" {
		q = q.Where("date >= ?", from)
	}
	if to != "" {
		q = q.Where("date <= ?", to)
	}

	switch group {
	case "model":
		q = q.Joins("JOIN models ON models.id = usage_dailies.model_id").
			Select("models.name as g, COALESCE(SUM(requests),0) as requests, COALESCE(SUM(input_tokens),0) as input_tokens, COALESCE(SUM(output_tokens),0) as output_tokens, COALESCE(SUM(cost),0) as cost").
			Group("models.name")
	default:
		q = q.Select("date as g, COALESCE(SUM(requests),0) as requests, COALESCE(SUM(input_tokens),0) as input_tokens, COALESCE(SUM(output_tokens),0) as output_tokens, COALESCE(SUM(cost),0) as cost").
			Group("date").Order("date ASC")
	}
	if err := q.Scan(&rows).Error; err != nil {
		h.log.Warn("usage stats failed", zap.Error(err))
		httpx.Fail(c, http.StatusInternalServerError, "usage stats failed")
		return
	}
	out := make([]gin.H, 0, len(rows))
	for _, r := range rows {
		out = append(out, gin.H{
			"key":          r.Key,
			"requests":     r.Requests,
			"inputTokens":  r.InputTokens,
			"outputTokens": r.OutputTokens,
			"cost":         r.Cost.String(),
		})
	}
	httpx.OK(c, out)
}

// Logs handles GET /api/v1/usage/logs?from&to&page&model&keyId.
func (h *Handler) Logs(c *gin.Context) {
	u := user.Current(c)
	page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
	pageSize := 20
	if page < 1 {
		page = 1
	}
	from := c.Query("from")
	to := c.Query("to")

	q := h.db.Model(&store.UsageLog{}).Where("user_id = ?", u.ID)
	if from != "" {
		q = q.Where("created_at >= ?", from)
	}
	if to != "" {
		q = q.Where("created_at <= ?", to+" 23:59:59")
	}
	if m := c.Query("model"); m != "" {
		q = q.Where("model_name = ?", m)
	}
	if kid := c.Query("keyId"); kid != "" {
		q = q.Where("key_id = ?", kid)
	}

	var total int64
	q.Count(&total)

	var logs []store.UsageLog
	if err := q.Order("id DESC").Offset((page - 1) * pageSize).Limit(pageSize).Find(&logs).Error; err != nil {
		h.log.Warn("usage logs failed", zap.Error(err))
		httpx.Fail(c, http.StatusInternalServerError, "usage logs failed")
		return
	}
	out := make([]gin.H, 0, len(logs))
	for _, l := range logs {
		out = append(out, gin.H{
			"id":               l.ID,
			"requestId":        l.RequestID,
			"model":            l.ModelName,
			"channelId":        l.ChannelID,
			"inputTokens":      l.InputTokens,
			"outputTokens":     l.OutputTokens,
			"cacheReadTokens":  l.CacheReadTokens,
			"cost":             l.Cost.String(),
			"latencyMs":        l.LatencyMs,
			"status":           l.Status,
			"errorCode":        l.ErrorCode,
			"createdAt":        l.CreatedAt,
		})
	}
	httpx.OK(c, gin.H{"total": total, "page": page, "pageSize": pageSize, "items": out})
}