package handlers import ( "net/http" "time" "github.com/gin-gonic/gin" "github.com/google/uuid" "gorm.io/gorm" "github.com/llamalink/llamalink/internal/api/middleware" "github.com/llamalink/llamalink/internal/db" "github.com/llamalink/llamalink/internal/quota" ) type UsageHandler struct { db *gorm.DB quotaSvc *quota.Service } func NewUsageHandler(db *gorm.DB, quotaSvc *quota.Service) *UsageHandler { return &UsageHandler{db: db, quotaSvc: quotaSvc} } func (h *UsageHandler) GetUsage(c *gin.Context) { keyIDStr := c.Param("key_id") keyID, err := uuid.Parse(keyIDStr) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": "invalid key id"}) return } // Only admins or key owner can view usage currentKey := middleware.GetAPIKey(c) if !currentKey.IsAdmin && currentKey.ID != keyID { c.JSON(http.StatusForbidden, gin.H{"error": "access denied"}) return } period := c.DefaultQuery("period", "month") now := time.Now().UTC() var periodStart, periodEnd time.Time switch period { case "month": periodStart = time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) periodEnd = periodStart.AddDate(0, 1, 0) case "year": periodStart = time.Date(now.Year(), 1, 1, 0, 0, 0, 0, time.UTC) periodEnd = periodStart.AddDate(1, 0, 0) default: periodStart = now.AddDate(0, 0, -7) periodEnd = now } // Get usage logs var logs []db.UsageLog h.db.Where("api_key_id = ? AND created_at >= ? AND created_at < ?", keyID, periodStart, periodEnd). Order("created_at DESC").Limit(100).Find(&logs) // Get aggregated stats var stats struct { TotalRequests int64 TotalTokens int64 AvgLatency float64 } h.db.Model(&db.UsageLog{}). Where("api_key_id = ? AND created_at >= ? AND created_at < ?", keyID, periodStart, periodEnd). Select("COUNT(*) as total_requests, COALESCE(SUM(total_tokens), 0) as total_tokens, COALESCE(AVG(latency_ms), 0) as avg_latency"). Scan(&stats) // Get quota info used, limit, _ := h.quotaSvc.GetUsage(keyID, periodStart, periodEnd) c.JSON(http.StatusOK, gin.H{ "period": gin.H{ "start": periodStart, "end": periodEnd, }, "usage": gin.H{ "total_requests": stats.TotalRequests, "total_tokens": stats.TotalTokens, "avg_latency_ms": stats.AvgLatency, }, "quota": gin.H{ "tokens_used": used, "tokens_limit": limit, }, "logs": logs, }) } func (h *UsageHandler) GetCurrentKeyUsage(c *gin.Context) { key := middleware.GetAPIKey(c) now := time.Now().UTC() periodStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) periodEnd := periodStart.AddDate(0, 1, 0) var stats struct { TotalRequests int64 TotalTokens int64 AvgLatency float64 } h.db.Model(&db.UsageLog{}). Where("api_key_id = ? AND created_at >= ? AND created_at < ?", key.ID, periodStart, periodEnd). Select("COUNT(*) as total_requests, COALESCE(SUM(total_tokens), 0) as total_tokens, COALESCE(AVG(latency_ms), 0) as avg_latency"). Scan(&stats) used, limit, _ := h.quotaSvc.GetUsage(key.ID, periodStart, periodEnd) c.JSON(http.StatusOK, gin.H{ "period": gin.H{ "start": periodStart, "end": periodEnd, }, "usage": gin.H{ "total_requests": stats.TotalRequests, "total_tokens": stats.TotalTokens, "avg_latency_ms": stats.AvgLatency, }, "quota": gin.H{ "tokens_used": used, "tokens_limit": limit, }, }) }