ratelimit.go

package tok

import (
	"fmt"
	"math"
	"strings"
	"sync"
	"time"
)

// UsageEntry represents a single recorded usage event.
type UsageEntry struct {
	Tokens    int
	CostUSD   float64
	Timestamp time.Time
	Provider  string
	Model     string
}

// Alert represents a usage threshold alert.
type Alert struct {
	Level     string // "warning", "critical", "limit_reached"
	Message   string
	Timestamp time.Time
	Threshold float64 // what % triggered it
}

// UsageSummary provides a snapshot of current usage across all windows.
type UsageSummary struct {
	HourlyTokens     int
	HourlyRemaining  int
	DailyTokens      int
	DailyRemaining   int
	SessionTokens    int
	SessionRemaining int
	DailyCostUSD     float64
	CostRemaining    float64
	HourlyPct        float64
	DailyPct         float64
}

// UsageTracker tracks API usage across sessions and prevents surprise bills.
type UsageTracker struct {
	DailyLimit   int
	HourlyLimit  int
	SessionLimit int
	CostLimitUSD float64

	hourlyUsage  []UsageEntry
	dailyUsage   []UsageEntry
	sessionUsage int
	mu           sync.Mutex
	Alerts       []Alert

	// Track which thresholds have already fired to avoid duplicates.
	// Keys are like "hourly_50", "daily_80", "cost_100", etc.
	firedThresholds map[string]bool
}

// NewUsageTracker creates a UsageTracker with sensible defaults.
func NewUsageTracker() *UsageTracker {
	return &UsageTracker{
		DailyLimit:      1_000_000,
		HourlyLimit:     200_000,
		SessionLimit:    500_000,
		CostLimitUSD:    10.00,
		hourlyUsage:     make([]UsageEntry, 0),
		dailyUsage:      make([]UsageEntry, 0),
		firedThresholds: make(map[string]bool),
	}
}

// Record adds a usage entry and checks thresholds.
func (u *UsageTracker) Record(tokens int, costUSD float64, provider, model string) {
	u.mu.Lock()
	defer u.mu.Unlock()

	entry := UsageEntry{
		Tokens:    tokens,
		CostUSD:   costUSD,
		Timestamp: time.Now(),
		Provider:  provider,
		Model:     model,
	}

	u.hourlyUsage = append(u.hourlyUsage, entry)
	u.dailyUsage = append(u.dailyUsage, entry)
	u.sessionUsage += tokens

	u.checkThresholdsLocked()
}

// CanProceed checks all limits and returns whether another request is allowed.
// If not, it returns false and a reason string.
func (u *UsageTracker) CanProceed() (bool, string) {
	u.mu.Lock()
	defer u.mu.Unlock()

	u.pruneOldLocked()

	hourlyTokens := u.hourlyTokensLocked()
	if hourlyTokens >= u.HourlyLimit {
		return false, fmt.Sprintf("hourly token limit reached (%d/%d)", hourlyTokens, u.HourlyLimit)
	}

	dailyTokens := u.dailyTokensLocked()
	if dailyTokens >= u.DailyLimit {
		return false, fmt.Sprintf("daily token limit reached (%d/%d)", dailyTokens, u.DailyLimit)
	}

	if u.sessionUsage >= u.SessionLimit {
		return false, fmt.Sprintf("session token limit reached (%d/%d)", u.sessionUsage, u.SessionLimit)
	}

	dailyCost := u.dailyCostLocked()
	if dailyCost >= u.CostLimitUSD {
		return false, fmt.Sprintf("daily cost limit reached ($%.2f/$%.2f)", dailyCost, u.CostLimitUSD)
	}

	return true, ""
}

// GetUsage returns a snapshot of current usage.
func (u *UsageTracker) GetUsage() UsageSummary {
	u.mu.Lock()
	defer u.mu.Unlock()

	u.pruneOldLocked()

	hourlyTokens := u.hourlyTokensLocked()
	dailyTokens := u.dailyTokensLocked()
	dailyCost := u.dailyCostLocked()

	hourlyRemaining := u.HourlyLimit - hourlyTokens
	if hourlyRemaining < 0 {
		hourlyRemaining = 0
	}
	dailyRemaining := u.DailyLimit - dailyTokens
	if dailyRemaining < 0 {
		dailyRemaining = 0
	}
	sessionRemaining := u.SessionLimit - u.sessionUsage
	if sessionRemaining < 0 {
		sessionRemaining = 0
	}
	costRemaining := u.CostLimitUSD - dailyCost
	if costRemaining < 0 {
		costRemaining = 0
	}

	var hourlyPct, dailyPct float64
	if u.HourlyLimit > 0 {
		hourlyPct = float64(hourlyTokens) / float64(u.HourlyLimit) * 100
	}
	if u.DailyLimit > 0 {
		dailyPct = float64(dailyTokens) / float64(u.DailyLimit) * 100
	}

	return UsageSummary{
		HourlyTokens:     hourlyTokens,
		HourlyRemaining:  hourlyRemaining,
		DailyTokens:      dailyTokens,
		DailyRemaining:   dailyRemaining,
		SessionTokens:    u.sessionUsage,
		SessionRemaining: sessionRemaining,
		DailyCostUSD:     dailyCost,
		CostRemaining:    costRemaining,
		HourlyPct:        hourlyPct,
		DailyPct:         dailyPct,
	}
}

// CheckThresholds evaluates current usage against threshold levels and generates alerts.
func (u *UsageTracker) CheckThresholds() {
	u.mu.Lock()
	defer u.mu.Unlock()

	u.checkThresholdsLocked()
}

func (u *UsageTracker) checkThresholdsLocked() {
	u.pruneOldLocked()

	hourlyTokens := u.hourlyTokensLocked()
	dailyTokens := u.dailyTokensLocked()
	dailyCost := u.dailyCostLocked()

	// Hourly thresholds
	if u.HourlyLimit > 0 {
		pct := float64(hourlyTokens) / float64(u.HourlyLimit) * 100
		u.emitAlert("hourly", pct, "hourly token usage")
	}

	// Daily thresholds
	if u.DailyLimit > 0 {
		pct := float64(dailyTokens) / float64(u.DailyLimit) * 100
		u.emitAlert("daily", pct, "daily token usage")
	}

	// Session thresholds
	if u.SessionLimit > 0 {
		pct := float64(u.sessionUsage) / float64(u.SessionLimit) * 100
		u.emitAlert("session", pct, "session token usage")
	}

	// Cost thresholds
	if u.CostLimitUSD > 0 {
		pct := dailyCost / u.CostLimitUSD * 100
		u.emitAlert("cost", pct, "daily cost")
	}
}

func (u *UsageTracker) emitAlert(category string, pct float64, label string) {
	type threshold struct {
		pct   float64
		level string
	}
	thresholds := []threshold{
		{100, "limit_reached"},
		{80, "critical"},
		{50, "warning"},
	}

	for _, t := range thresholds {
		if pct >= t.pct {
			key := fmt.Sprintf("%s_%d", category, int(t.pct))
			if !u.firedThresholds[key] {
				u.firedThresholds[key] = true
				u.Alerts = append(u.Alerts, Alert{
					Level:     t.level,
					Message:   fmt.Sprintf("%s at %.0f%% of limit", label, pct),
					Timestamp: time.Now(),
					Threshold: t.pct,
				})
			}
			// Only fire the highest applicable threshold
			return
		}
	}
}

// Reset clears the session counter and alerts.
func (u *UsageTracker) Reset() {
	u.mu.Lock()
	defer u.mu.Unlock()

	u.sessionUsage = 0
	u.Alerts = nil
	u.firedThresholds = make(map[string]bool)
}

// PruneOld removes entries older than their respective windows.
func (u *UsageTracker) PruneOld() {
	u.mu.Lock()
	defer u.mu.Unlock()

	u.pruneOldLocked()
}

func (u *UsageTracker) pruneOldLocked() {
	now := time.Now()
	hourAgo := now.Add(-1 * time.Hour)
	dayAgo := now.Add(-24 * time.Hour)

	// Prune hourly
	pruned := u.hourlyUsage[:0]
	for _, e := range u.hourlyUsage {
		if !e.Timestamp.Before(hourAgo) {
			pruned = append(pruned, e)
		}
	}
	u.hourlyUsage = pruned

	// Prune daily
	pruned = u.dailyUsage[:0]
	for _, e := range u.dailyUsage {
		if !e.Timestamp.Before(dayAgo) {
			pruned = append(pruned, e)
		}
	}
	u.dailyUsage = pruned
}

// EstimateRemaining estimates how many more requests of the given size fit in the budget.
func (u *UsageTracker) EstimateRemaining(tokensPerRequest int) int {
	u.mu.Lock()
	defer u.mu.Unlock()

	if tokensPerRequest <= 0 {
		return 0
	}

	u.pruneOldLocked()

	hourlyTokens := u.hourlyTokensLocked()
	dailyTokens := u.dailyTokensLocked()

	hourlyRemaining := u.HourlyLimit - hourlyTokens
	dailyRemaining := u.DailyLimit - dailyTokens
	sessionRemaining := u.SessionLimit - u.sessionUsage

	// Find the minimum remaining across all token-based limits
	minRemaining := hourlyRemaining
	if dailyRemaining < minRemaining {
		minRemaining = dailyRemaining
	}
	if sessionRemaining < minRemaining {
		minRemaining = sessionRemaining
	}

	if minRemaining <= 0 {
		return 0
	}

	return minRemaining / tokensPerRequest
}

// FormatUsageBar creates an ASCII progress bar.
func FormatUsageBar(pct float64, width int) string {
	if width <= 0 {
		return ""
	}
	if pct < 0 {
		pct = 0
	}
	if pct > 100 {
		pct = 100
	}

	filled := int(math.Round(pct / 100 * float64(width)))
	if filled > width {
		filled = width
	}

	bar := strings.Repeat("█", filled) + strings.Repeat("░", width-filled)
	return fmt.Sprintf("[%s] %d%%", bar, int(pct))
}

// FormatSummary returns a human-readable usage summary.
func (u *UsageTracker) FormatSummary() string {
	summary := u.GetUsage()

	sessionPct := float64(0)
	if u.SessionLimit > 0 {
		sessionPct = float64(summary.SessionTokens) / float64(u.SessionLimit) * 100
	}
	costPct := float64(0)
	if u.CostLimitUSD > 0 {
		costPct = summary.DailyCostUSD / u.CostLimitUSD * 100
	}

	barWidth := 16

	var sb strings.Builder
	sb.WriteString("Token Usage:\n")
	fmt.Fprintf(&sb, "  Hourly:  %s / %s (%d%%) %s\n",
		formatNumber(summary.HourlyTokens),
		formatNumber(u.HourlyLimit),
		int(summary.HourlyPct),
		FormatUsageBar(summary.HourlyPct, barWidth))
	fmt.Fprintf(&sb, "  Daily:  %s / %s (%d%%) %s\n",
		formatNumber(summary.DailyTokens),
		formatNumber(u.DailyLimit),
		int(summary.DailyPct),
		FormatUsageBar(summary.DailyPct, barWidth))
	fmt.Fprintf(&sb, "  Session: %s / %s (%d%%) %s\n",
		formatNumber(summary.SessionTokens),
		formatNumber(u.SessionLimit),
		int(sessionPct),
		FormatUsageBar(sessionPct, barWidth))
	fmt.Fprintf(&sb, "  Cost:   $%.2f / $%.2f (%d%%) %s",
		summary.DailyCostUSD,
		u.CostLimitUSD,
		int(costPct),
		FormatUsageBar(costPct, barWidth))

	return sb.String()
}

// --- internal helpers ---

func (u *UsageTracker) hourlyTokensLocked() int {
	total := 0
	for _, e := range u.hourlyUsage {
		total += e.Tokens
	}
	return total
}

func (u *UsageTracker) dailyTokensLocked() int {
	total := 0
	for _, e := range u.dailyUsage {
		total += e.Tokens
	}
	return total
}

func (u *UsageTracker) dailyCostLocked() float64 {
	total := 0.0
	for _, e := range u.dailyUsage {
		total += e.CostUSD
	}
	return total
}

// formatNumber is defined in format.go and shared across the package.