package firewall

import (
	"net/netip"
	"sync/atomic"
	"testing"
	"time"
)

func TestRateLimiterAllow(t *testing.T) {
	rl := newRateLimiter(3, time.Minute)
	defer rl.Stop()

	addr := netip.MustParseAddr("10.0.0.1")

	for i := 0; i < 3; i++ {
		if !rl.Allow(addr) {
			t.Fatalf("call %d: expected Allow=true", i+1)
		}
	}

	if rl.Allow(addr) {
		t.Fatal("call 4: expected Allow=false (over limit)")
	}
}

func TestRateLimiterDifferentIPs(t *testing.T) {
	rl := newRateLimiter(2, time.Minute)
	defer rl.Stop()

	a := netip.MustParseAddr("10.0.0.1")
	b := netip.MustParseAddr("10.0.0.2")

	// Exhaust a's limit.
	for i := 0; i < 2; i++ {
		rl.Allow(a)
	}
	if rl.Allow(a) {
		t.Fatal("a should be rate limited")
	}

	// b should be independent.
	if !rl.Allow(b) {
		t.Fatal("b should not be rate limited")
	}
}

func TestRateLimiterWindowReset(t *testing.T) {
	rl := newRateLimiter(2, time.Minute)
	defer rl.Stop()

	var fakeNow atomic.Int64
	fakeNow.Store(time.Now().UnixNano())
	rl.now = func() time.Time {
		return time.Unix(0, fakeNow.Load())
	}

	addr := netip.MustParseAddr("10.0.0.1")

	// Exhaust the limit.
	rl.Allow(addr)
	rl.Allow(addr)
	if rl.Allow(addr) {
		t.Fatal("should be rate limited")
	}

	// Advance past the window.
	fakeNow.Add(int64(2 * time.Minute))

	// Should be allowed again.
	if !rl.Allow(addr) {
		t.Fatal("should be allowed after window reset")
	}
}

func TestRateLimiterCleanup(t *testing.T) {
	rl := newRateLimiter(10, 50*time.Millisecond)
	defer rl.Stop()

	addr := netip.MustParseAddr("10.0.0.1")
	rl.Allow(addr)

	// Entry should exist.
	if _, ok := rl.entries.Load(addr); !ok {
		t.Fatal("entry should exist")
	}

	// Wait for 2*window + a cleanup cycle to pass.
	time.Sleep(200 * time.Millisecond)

	// Entry should have been cleaned up.
	if _, ok := rl.entries.Load(addr); ok {
		t.Fatal("stale entry should have been cleaned up")
	}
}