package api

import (
	"testing"
	"time"
)

func resetAgentWSRateLim(t *testing.T) {
	t.Helper()
	agentWSRateLim.mu.Lock()
	agentWSRateLim.attempts = make(map[string][]time.Time)
	agentWSRateLim.mu.Unlock()
}

func TestAllowAgentWSUpgradeRateLimit(t *testing.T) {
	t.Run("rejects 31st attempt within window", func(t *testing.T) {
		resetAgentWSRateLim(t)
		ip := "203.0.113.42"

		for i := 1; i <= agentWSRateLimitMax; i++ {
			if !allowAgentWSUpgrade(ip) {
				t.Fatalf("attempt %d: expected allow, got reject", i)
			}
		}
		if allowAgentWSUpgrade(ip) {
			t.Fatal("31st attempt: expected reject, got allow")
		}
		if allowAgentWSUpgrade(ip) {
			t.Fatal("32nd attempt: expected reject, got allow")
		}
	})

	t.Run("empty IP bypasses limit", func(t *testing.T) {
		resetAgentWSRateLim(t)

		for i := 1; i <= agentWSRateLimitMax+5; i++ {
			if !allowAgentWSUpgrade("") {
				t.Fatalf("empty IP attempt %d: expected allow, got reject", i)
			}
		}
	})

	t.Run("stale attempts outside window are pruned", func(t *testing.T) {
		resetAgentWSRateLim(t)
		ip := "198.51.100.7"
		stale := time.Now().Add(-agentWSRateLimitWindow - time.Second)

		agentWSRateLim.mu.Lock()
		staleAttempts := make([]time.Time, agentWSRateLimitMax)
		for i := range staleAttempts {
			staleAttempts[i] = stale
		}
		agentWSRateLim.attempts[ip] = staleAttempts
		agentWSRateLim.mu.Unlock()

		if !allowAgentWSUpgrade(ip) {
			t.Fatal("expected allow after stale attempts pruned")
		}
	})

	t.Run("different IPs have independent limits", func(t *testing.T) {
		resetAgentWSRateLim(t)
		ipA := "192.0.2.1"
		ipB := "192.0.2.2"

		for i := 1; i <= agentWSRateLimitMax; i++ {
			if !allowAgentWSUpgrade(ipA) {
				t.Fatalf("ipA attempt %d: expected allow, got reject", i)
			}
		}
		if !allowAgentWSUpgrade(ipB) {
			t.Fatal("ipB first attempt: expected allow after ipA exhausted")
		}
	})
}