ratelimit_test.go

package tok

import (
	"strings"
	"sync"
	"testing"
	"time"
)

func TestRecordIncreasesCounters(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.Record(1000, 0.05, "openai", "gpt-4")

	summary := tracker.GetUsage()
	if summary.HourlyTokens != 1000 {
		t.Errorf("expected hourly tokens 1000, got %d", summary.HourlyTokens)
	}
	if summary.DailyTokens != 1000 {
		t.Errorf("expected daily tokens 1000, got %d", summary.DailyTokens)
	}
	if summary.SessionTokens != 1000 {
		t.Errorf("expected session tokens 1000, got %d", summary.SessionTokens)
	}
	if summary.DailyCostUSD != 0.05 {
		t.Errorf("expected daily cost 0.05, got %f", summary.DailyCostUSD)
	}
}

func TestCanProceedReturnsFalseWhenHourlyLimitHit(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.HourlyLimit = 1000

	tracker.Record(1000, 0.01, "openai", "gpt-4")

	ok, reason := tracker.CanProceed()
	if ok {
		t.Error("expected CanProceed to return false when hourly limit reached")
	}
	if !strings.Contains(reason, "hourly") {
		t.Errorf("expected reason to mention hourly, got: %s", reason)
	}
}

func TestCanProceedReturnsFalseWhenDailyLimitHit(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.DailyLimit = 5000
	tracker.HourlyLimit = 100000 // high so it doesn't trigger first

	tracker.Record(5000, 0.01, "openai", "gpt-4")

	ok, reason := tracker.CanProceed()
	if ok {
		t.Error("expected CanProceed to return false when daily limit reached")
	}
	if !strings.Contains(reason, "daily token") {
		t.Errorf("expected reason to mention daily token, got: %s", reason)
	}
}

func TestCanProceedReturnsFalseWhenSessionLimitHit(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.SessionLimit = 2000
	tracker.HourlyLimit = 100000
	tracker.DailyLimit = 100000

	tracker.Record(2000, 0.01, "openai", "gpt-4")

	ok, reason := tracker.CanProceed()
	if ok {
		t.Error("expected CanProceed to return false when session limit reached")
	}
	if !strings.Contains(reason, "session") {
		t.Errorf("expected reason to mention session, got: %s", reason)
	}
}

func TestCanProceedReturnsFalseWhenCostLimitHit(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.CostLimitUSD = 1.00
	tracker.HourlyLimit = 10000000
	tracker.DailyLimit = 10000000
	tracker.SessionLimit = 10000000

	tracker.Record(100, 1.00, "openai", "gpt-4")

	ok, reason := tracker.CanProceed()
	if ok {
		t.Error("expected CanProceed to return false when cost limit reached")
	}
	if !strings.Contains(reason, "cost") {
		t.Errorf("expected reason to mention cost, got: %s", reason)
	}
}

func TestHourlyWindowExpiry(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.HourlyLimit = 1000

	// Inject an old entry directly
	tracker.mu.Lock()
	tracker.hourlyUsage = append(tracker.hourlyUsage, UsageEntry{
		Tokens:    900,
		CostUSD:   0.01,
		Timestamp: time.Now().Add(-2 * time.Hour),
		Provider:  "openai",
		Model:     "gpt-4",
	})
	tracker.mu.Unlock()

	// After pruning, the old entry should be gone
	tracker.PruneOld()

	summary := tracker.GetUsage()
	if summary.HourlyTokens != 0 {
		t.Errorf("expected hourly tokens 0 after expiry, got %d", summary.HourlyTokens)
	}
}

func TestDailyWindowExpiry(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.DailyLimit = 5000

	// Inject an old entry directly
	tracker.mu.Lock()
	tracker.dailyUsage = append(tracker.dailyUsage, UsageEntry{
		Tokens:    4000,
		CostUSD:   0.50,
		Timestamp: time.Now().Add(-25 * time.Hour),
		Provider:  "openai",
		Model:     "gpt-4",
	})
	tracker.mu.Unlock()

	tracker.PruneOld()

	summary := tracker.GetUsage()
	if summary.DailyTokens != 0 {
		t.Errorf("expected daily tokens 0 after expiry, got %d", summary.DailyTokens)
	}
}

func TestAlertGenerationAtThresholds(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.HourlyLimit = 100

	// 50% - warning
	tracker.Record(50, 0.01, "openai", "gpt-4")

	found := false
	for _, a := range tracker.Alerts {
		if a.Level == "warning" && a.Threshold == 50 {
			found = true
		}
	}
	if !found {
		t.Error("expected warning alert at 50% threshold")
	}

	// 80% - critical
	tracker.Record(30, 0.01, "openai", "gpt-4")

	found = false
	for _, a := range tracker.Alerts {
		if a.Level == "critical" && a.Threshold == 80 {
			found = true
		}
	}
	if !found {
		t.Error("expected critical alert at 80% threshold")
	}

	// 100% - limit_reached
	tracker.Record(20, 0.01, "openai", "gpt-4")

	found = false
	for _, a := range tracker.Alerts {
		if a.Level == "limit_reached" && a.Threshold == 100 {
			found = true
		}
	}
	if !found {
		t.Error("expected limit_reached alert at 100% threshold")
	}
}

func TestNoDuplicateAlerts(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.HourlyLimit = 100

	// Record multiple times past 50%
	tracker.Record(55, 0.01, "openai", "gpt-4")
	tracker.Record(5, 0.01, "openai", "gpt-4")
	tracker.Record(5, 0.01, "openai", "gpt-4")

	// Count how many warning-level alerts with threshold 50 exist for hourly
	count := 0
	for _, a := range tracker.Alerts {
		if a.Threshold == 50 && a.Level == "warning" && strings.Contains(a.Message, "hourly") {
			count++
		}
	}
	if count > 1 {
		t.Errorf("expected at most 1 warning alert for hourly 50%%, got %d", count)
	}
}

func TestFormatUsageBarEmpty(t *testing.T) {
	bar := FormatUsageBar(0, 16)
	if !strings.Contains(bar, "0%") {
		t.Errorf("expected 0%% in bar, got: %s", bar)
	}
	// Should have all empty blocks
	if !strings.Contains(bar, "░░░░░░░░░░░░░░░░") {
		t.Errorf("expected all empty blocks, got: %s", bar)
	}
}

func TestFormatUsageBarFull(t *testing.T) {
	bar := FormatUsageBar(100, 16)
	if !strings.Contains(bar, "100%") {
		t.Errorf("expected 100%% in bar, got: %s", bar)
	}
	if !strings.Contains(bar, "████████████████") {
		t.Errorf("expected all filled blocks, got: %s", bar)
	}
}

func TestFormatUsageBarHalf(t *testing.T) {
	bar := FormatUsageBar(50, 16)
	if !strings.Contains(bar, "50%") {
		t.Errorf("expected 50%% in bar, got: %s", bar)
	}
}

func TestFormatSummaryOutput(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.Record(45000, 2.40, "openai", "gpt-4")

	output := tracker.FormatSummary()

	if !strings.Contains(output, "Token Usage:") {
		t.Error("expected output to contain 'Token Usage:'")
	}
	if !strings.Contains(output, "Hourly:") {
		t.Error("expected output to contain 'Hourly:'")
	}
	if !strings.Contains(output, "Daily:") {
		t.Error("expected output to contain 'Daily:'")
	}
	if !strings.Contains(output, "Session:") {
		t.Error("expected output to contain 'Session:'")
	}
	if !strings.Contains(output, "Cost:") {
		t.Error("expected output to contain 'Cost:'")
	}
	if !strings.Contains(output, "$2.40") {
		t.Errorf("expected output to contain '$2.40', got:\n%s", output)
	}
}

func TestEstimateRemaining(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.HourlyLimit = 10000
	tracker.DailyLimit = 100000
	tracker.SessionLimit = 50000

	tracker.Record(5000, 0.10, "openai", "gpt-4")

	// Hourly remaining: 10000-5000=5000, daily: 100000-5000=95000, session: 50000-5000=45000
	// Min remaining is 5000 (hourly), so 5000/1000 = 5
	est := tracker.EstimateRemaining(1000)
	if est != 5 {
		t.Errorf("expected estimate 5, got %d", est)
	}
}

func TestEstimateRemainingZeroTokens(t *testing.T) {
	tracker := NewUsageTracker()
	est := tracker.EstimateRemaining(0)
	if est != 0 {
		t.Errorf("expected 0 for zero tokensPerRequest, got %d", est)
	}
}

func TestConcurrentRecordCalls(t *testing.T) {
	tracker := NewUsageTracker()

	var wg sync.WaitGroup
	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			tracker.Record(100, 0.01, "openai", "gpt-4")
		}()
	}
	wg.Wait()

	summary := tracker.GetUsage()
	if summary.SessionTokens != 10000 {
		t.Errorf("expected session tokens 10000 after concurrent records, got %d", summary.SessionTokens)
	}
}

func TestResetClearsSession(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.Record(5000, 0.50, "openai", "gpt-4")

	tracker.Reset()

	summary := tracker.GetUsage()
	if summary.SessionTokens != 0 {
		t.Errorf("expected session tokens 0 after reset, got %d", summary.SessionTokens)
	}

	if len(tracker.Alerts) != 0 {
		t.Errorf("expected alerts cleared after reset, got %d", len(tracker.Alerts))
	}
}

func TestCanProceedReturnsTrueUnderLimits(t *testing.T) {
	tracker := NewUsageTracker()
	tracker.Record(100, 0.01, "openai", "gpt-4")

	ok, reason := tracker.CanProceed()
	if !ok {
		t.Errorf("expected CanProceed to return true under limits, got reason: %s", reason)
	}
}

func TestFormatNumber(t *testing.T) {
	tests := []struct {
		input    int
		expected string
	}{
		{0, "0"},
		{999, "999"},
		{1000, "1,000"},
		{45000, "45,000"},
		{1000000, "1,000,000"},
		{200000, "200,000"},
	}

	for _, tc := range tests {
		got := formatNumber(tc.input)
		if got != tc.expected {
			t.Errorf("formatNumber(%d) = %s, want %s", tc.input, got, tc.expected)
		}
	}
}