package cron import ( "context" "log/slog" "strings" "time" "github.com/ez-api/ez-api/internal/model" "github.com/ez-api/ez-api/internal/service" "gorm.io/gorm" ) type QuotaResetter struct { db *gorm.DB sync *service.SyncService } func NewQuotaResetter(db *gorm.DB, sync *service.SyncService) *QuotaResetter { return &QuotaResetter{db: db, sync: sync} } // RunOnce executes a single quota reset check. Called by scheduler. func (q *QuotaResetter) RunOnce(ctx context.Context) { if q == nil || q.db == nil { return } if err := q.resetOnce(ctx); err != nil { slog.Default().Warn("quota reset failed", "err", err) } } func (q *QuotaResetter) resetOnce(ctx context.Context) error { if q == nil || q.db == nil { return nil } now := time.Now().UTC() var keys []model.Key if err := q.db.Where("quota_reset_type IN ? AND (quota_reset_at IS NULL OR quota_reset_at <= ?)", []string{"daily", "monthly"}, now).Find(&keys).Error; err != nil { return err } for i := range keys { resetType := strings.ToLower(strings.TrimSpace(keys[i].QuotaResetType)) nextAt, ok := nextQuotaReset(now, resetType) if !ok { continue } if err := q.db.Model(&keys[i]).Updates(map[string]any{ "quota_used": 0, "quota_reset_at": nextAt, }).Error; err != nil { slog.Default().Warn("quota reset update failed", "key_id", keys[i].ID, "err", err) continue } keys[i].QuotaUsed = 0 keys[i].QuotaResetAt = &nextAt if q.sync != nil { _ = q.sync.SyncKey(&keys[i]) } } return nil } func nextQuotaReset(now time.Time, resetType string) (time.Time, bool) { now = now.UTC() switch resetType { case "daily": next := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).AddDate(0, 0, 1) return next, true case "monthly": next := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC).AddDate(0, 1, 0) return next, true default: return time.Time{}, false } }