llama-link/internal/api/handlers/usage.go

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,
		},
	})
}