ez-api/internal/api/stats_handler.go

package api

import (
	"fmt"
	"net/http"
	"sort"
	"strings"
	"time"

	"github.com/ez-api/ez-api/internal/model"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

type KeyUsageStat struct {
	KeyID    uint  `json:"key_id"`
	Requests int64 `json:"requests"`
	Tokens   int64 `json:"tokens"`
}

type ModelUsageStat struct {
	Model    string `json:"model"`
	Requests int64  `json:"requests"`
	Tokens   int64  `json:"tokens"`
}

type MasterUsageStatsResponse struct {
	Period        string           `json:"period,omitempty"`
	TotalRequests int64            `json:"total_requests"`
	TotalTokens   int64            `json:"total_tokens"`
	ByKey         []KeyUsageStat   `json:"by_key"`
	ByModel       []ModelUsageStat `json:"by_model"`
}

// GetSelfStats godoc
// @Summary      Usage stats (master)
// @Description  Aggregate request stats for the authenticated master
// @Tags         master
// @Produce      json
// @Security     MasterAuth
// @Param        period  query     string  false  "today|week|month|all"
// @Param        since   query     int     false  "unix seconds"
// @Param        until   query     int     false  "unix seconds"
// @Success      200  {object}  ResponseEnvelope{data=MasterUsageStatsResponse}
// @Failure      400  {object}  ResponseEnvelope{data=gin.H}
// @Failure      401  {object}  ResponseEnvelope{data=gin.H}
// @Failure      500  {object}  ResponseEnvelope{data=gin.H}
// @Router       /v1/stats [get]
func (h *MasterHandler) GetSelfStats(c *gin.Context) {
	master, exists := c.Get("master")
	if !exists {
		c.JSON(http.StatusUnauthorized, gin.H{"error": "master key not found in context"})
		return
	}
	m := master.(*model.Master)

	rng, err := parseStatsRange(c)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	logDB := h.logDBConn()
	base := h.logBaseQuery()
	if logDB == h.db {
		base = base.Joins("JOIN keys ON keys.id = log_records.key_id").
			Where("keys.master_id = ?", m.ID)
	} else {
		var keyIDs []uint
		if err := h.db.Model(&model.Key{}).
			Where("master_id = ?", m.ID).
			Pluck("id", &keyIDs).Error; err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load master keys", "details": err.Error()})
			return
		}
		if len(keyIDs) == 0 {
			base = base.Where("1 = 0")
		} else {
			base = base.Where("log_records.key_id IN ?", keyIDs)
		}
	}
	base = applyStatsRange(base, rng)

	totalRequests, totalTokens, err := aggregateTotals(base)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to aggregate stats", "details": err.Error()})
		return
	}

	var byKey []KeyUsageStat
	if err := base.Session(&gorm.Session{}).
		Select("log_records.key_id as key_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
		Group("log_records.key_id").
		Scan(&byKey).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by key", "details": err.Error()})
		return
	}

	var byModel []ModelUsageStat
	if err := base.Session(&gorm.Session{}).
		Select("log_records.model_name as model, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
		Group("log_records.model_name").
		Scan(&byModel).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by model", "details": err.Error()})
		return
	}

	c.JSON(http.StatusOK, MasterUsageStatsResponse{
		Period:        rng.Period,
		TotalRequests: totalRequests,
		TotalTokens:   totalTokens,
		ByKey:         byKey,
		ByModel:       byModel,
	})
}

type MasterUsageAgg struct {
	MasterID uint  `json:"master_id"`
	Requests int64 `json:"requests"`
	Tokens   int64 `json:"tokens"`
}

type ProviderUsageAgg struct {
	ProviderID   uint   `json:"provider_id"`
	ProviderType string `json:"provider_type"`
	ProviderName string `json:"provider_name"`
	Requests     int64  `json:"requests"`
	Tokens       int64  `json:"tokens"`
}

type AdminUsageStatsResponse struct {
	Period        string             `json:"period,omitempty"`
	TotalMasters  int64              `json:"total_masters"`
	ActiveMasters int64              `json:"active_masters"`
	TotalRequests int64              `json:"total_requests"`
	TotalTokens   int64              `json:"total_tokens"`
	ByMaster      []MasterUsageAgg   `json:"by_master"`
	ByProvider    []ProviderUsageAgg `json:"by_provider"`
}

// GetAdminStats godoc
// @Summary      Usage stats (admin)
// @Description  Aggregate request stats across all masters
// @Tags         admin
// @Produce      json
// @Security     AdminAuth
// @Param        period  query     string  false  "today|week|month|all"
// @Param        since   query     int     false  "unix seconds"
// @Param        until   query     int     false  "unix seconds"
// @Success      200  {object}  ResponseEnvelope{data=AdminUsageStatsResponse}
// @Failure      400  {object}  ResponseEnvelope{data=gin.H}
// @Failure      500  {object}  ResponseEnvelope{data=gin.H}
// @Router       /admin/stats [get]
func (h *AdminHandler) GetAdminStats(c *gin.Context) {
	rng, err := parseStatsRange(c)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	var totalMasters int64
	if err := h.db.Model(&model.Master{}).Count(&totalMasters).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to count masters", "details": err.Error()})
		return
	}
	var activeMasters int64
	if err := h.db.Model(&model.Master{}).Where("status = ?", "active").Count(&activeMasters).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to count active masters", "details": err.Error()})
		return
	}

	logDB := h.logDBConn()
	base := h.logBaseQuery()
	base = applyStatsRange(base, rng)

	totalRequests, totalTokens, err := aggregateTotals(base)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to aggregate stats", "details": err.Error()})
		return
	}

	var byMaster []MasterUsageAgg
	if logDB == h.db {
		if err := base.Session(&gorm.Session{}).
			Joins("JOIN keys ON keys.id = log_records.key_id").
			Select("keys.master_id as master_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
			Group("keys.master_id").
			Scan(&byMaster).Error; err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by master", "details": err.Error()})
			return
		}
	} else {
		var err error
		byMaster, err = aggregateByMasterFromLogs(base, h.db)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by master", "details": err.Error()})
			return
		}
	}

	var byProvider []ProviderUsageAgg
	if err := base.Session(&gorm.Session{}).
		Select("log_records.provider_id as provider_id, log_records.provider_type as provider_type, log_records.provider_name as provider_name, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
		Group("log_records.provider_id, log_records.provider_type, log_records.provider_name").
		Scan(&byProvider).Error; err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to group by provider", "details": err.Error()})
		return
	}

	c.JSON(http.StatusOK, AdminUsageStatsResponse{
		Period:        rng.Period,
		TotalMasters:  totalMasters,
		ActiveMasters: activeMasters,
		TotalRequests: totalRequests,
		TotalTokens:   totalTokens,
		ByMaster:      byMaster,
		ByProvider:    byProvider,
	})
}

type statsRange struct {
	Since  *time.Time
	Until  *time.Time
	Period string
}

func parseStatsRange(c *gin.Context) (statsRange, error) {
	period := strings.ToLower(strings.TrimSpace(c.Query("period")))
	if period != "" {
		if period == "all" {
			return statsRange{Period: period}, nil
		}
		start, now := periodWindow(period)
		if start.IsZero() {
			return statsRange{}, fmt.Errorf("invalid period")
		}
		return statsRange{Since: &start, Until: &now, Period: period}, nil
	}

	var since *time.Time
	if t, ok := parseUnixSeconds(c.Query("since")); ok {
		since = &t
	}
	var until *time.Time
	if t, ok := parseUnixSeconds(c.Query("until")); ok {
		until = &t
	}
	return statsRange{Since: since, Until: until}, nil
}

func periodWindow(period string) (time.Time, time.Time) {
	now := time.Now().UTC()
	startOfDay := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC)
	switch period {
	case "today":
		return startOfDay, now
	case "week":
		weekday := int(startOfDay.Weekday())
		if weekday == 0 {
			weekday = 7
		}
		start := startOfDay.AddDate(0, 0, -(weekday - 1))
		return start, now
	case "month":
		start := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
		return start, now
	case "last7d":
		start := now.AddDate(0, 0, -7)
		return start, now
	case "last30d":
		start := now.AddDate(0, 0, -30)
		return start, now
	default:
		return time.Time{}, time.Time{}
	}
}

func applyStatsRange(q *gorm.DB, rng statsRange) *gorm.DB {
	if rng.Since != nil {
		q = q.Where("log_records.created_at >= ?", *rng.Since)
	}
	if rng.Until != nil {
		q = q.Where("log_records.created_at <= ?", *rng.Until)
	}
	return q
}

func aggregateTotals(q *gorm.DB) (int64, int64, error) {
	var totalRequests int64
	if err := q.Session(&gorm.Session{}).Count(&totalRequests).Error; err != nil {
		return 0, 0, err
	}
	type totals struct {
		Tokens int64
	}
	var t totals
	if err := q.Session(&gorm.Session{}).
		Select("COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
		Scan(&t).Error; err != nil {
		return 0, 0, err
	}
	return totalRequests, t.Tokens, nil
}

func aggregateByMasterFromLogs(base *gorm.DB, mainDB *gorm.DB) ([]MasterUsageAgg, error) {
	if base == nil || mainDB == nil {
		return nil, nil
	}
	var byKey []KeyUsageStat
	if err := base.Session(&gorm.Session{}).
		Select("log_records.key_id as key_id, COUNT(*) as requests, COALESCE(SUM(log_records.tokens_in + log_records.tokens_out),0) as tokens").
		Group("log_records.key_id").
		Scan(&byKey).Error; err != nil {
		return nil, err
	}
	if len(byKey) == 0 {
		return []MasterUsageAgg{}, nil
	}

	keyIDs := make([]uint, 0, len(byKey))
	for _, row := range byKey {
		if row.KeyID > 0 {
			keyIDs = append(keyIDs, row.KeyID)
		}
	}
	if len(keyIDs) == 0 {
		return []MasterUsageAgg{}, nil
	}

	type keyMaster struct {
		ID       uint
		MasterID uint
	}
	var keyMasters []keyMaster
	if err := mainDB.Model(&model.Key{}).
		Select("id, master_id").
		Where("id IN ?", keyIDs).
		Scan(&keyMasters).Error; err != nil {
		return nil, err
	}
	keyToMaster := make(map[uint]uint, len(keyMasters))
	for _, row := range keyMasters {
		keyToMaster[row.ID] = row.MasterID
	}

	agg := make(map[uint]*MasterUsageAgg)
	for _, row := range byKey {
		masterID := keyToMaster[row.KeyID]
		if masterID == 0 {
			continue
		}
		entry, ok := agg[masterID]
		if !ok {
			entry = &MasterUsageAgg{MasterID: masterID}
			agg[masterID] = entry
		}
		entry.Requests += row.Requests
		entry.Tokens += row.Tokens
	}

	out := make([]MasterUsageAgg, 0, len(agg))
	for _, row := range agg {
		out = append(out, *row)
	}
	sort.Slice(out, func(i, j int) bool {
		return out[i].MasterID < out[j].MasterID
	})
	return out, nil
}