// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0

package ratelimit

import (
	"fmt"
	"testing"
	"testing/synctest"
	"time"
)

func TestLimiterAllowsWithinBudget(t *testing.T) {
	l := New(3, time.Minute)
	for i := range 3 {
		allowed, remaining, retry := l.Check("ip")
		if !allowed || remaining != 2-i || retry != 0 {
			t.Fatalf("hit %d: allowed=%v remaining=%d retry=%d", i, allowed, remaining, retry)
		}
	}
}

func TestLimiterBlocksOverBudget(t *testing.T) {
	l := New(2, time.Minute)
	l.Check("ip")
	l.Check("ip")
	allowed, remaining, retry := l.Check("ip")
	if allowed {
		t.Fatal("third hit allowed")
	}
	if remaining != 0 {
		t.Fatalf("remaining = %d", remaining)
	}
	if retry < 1 || retry > 61 {
		t.Fatalf("retry_after = %d", retry)
	}
}

// synctest runs the window in virtual time, so the boundary is exact and
// the test does not spend a millisecond of the wall clock on it.
func TestLimiterWindowExpiry(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		l := New(1, 20*time.Millisecond)
		if allowed, _, _ := l.Check("ip"); !allowed {
			t.Fatal("first hit blocked")
		}
		if allowed, _, _ := l.Check("ip"); allowed {
			t.Fatal("second hit allowed")
		}
		time.Sleep(19 * time.Millisecond)
		if allowed, _, _ := l.Check("ip"); allowed {
			t.Fatal("a hit inside the window was allowed")
		}
		// Past the window measured from the last recorded hit.
		time.Sleep(21 * time.Millisecond)
		if allowed, remaining, _ := l.Check("ip"); !allowed || remaining != 0 {
			t.Fatalf("hit after window: allowed=%v remaining=%d", allowed, remaining)
		}
	})
}

func TestLimiterTracksKeysSeparately(t *testing.T) {
	l := New(1, time.Minute)
	l.Check("a")
	if allowed, _, _ := l.Check("b"); !allowed {
		t.Fatal("separate key blocked")
	}
}

// A limiter built for one budget does not carry another's hits: the
// process-level limiters are constructed once, and a fresh one is clean.
func TestLimiterStartsClean(t *testing.T) {
	l := New(1, time.Minute)
	if allowed, _, _ := l.Check("ip"); !allowed {
		t.Fatal("the first hit was refused")
	}
	if allowed, _, _ := l.Check("ip"); allowed {
		t.Fatal("the second hit was allowed over a budget of one")
	}
	fresh := New(10, 20*time.Millisecond)
	if allowed, _, _ := fresh.Check("ip"); !allowed {
		t.Fatal("a fresh limiter started with hits recorded")
	}
}

func TestLoginLimiterBlocksAfterMaxAttempts(t *testing.T) {
	l := NewLoginLimiter()
	for i := 1; i <= loginMaxAttempts; i++ {
		if count := l.Record("ip"); count != i {
			t.Fatalf("count = %d, want %d", count, i)
		}
		if blocked, _ := l.Blocked("ip"); blocked {
			t.Fatalf("blocked at attempt %d", i)
		}
	}
	if count := l.Record("ip"); count != loginMaxAttempts+1 {
		t.Fatalf("count = %d", count)
	}
	blocked, retry := l.Blocked("ip")
	if !blocked {
		t.Fatal("not blocked after exceeding the budget")
	}
	if retry < 1 || retry > 61 {
		t.Fatalf("retry_after = %d", retry)
	}
}

func TestLoginLimiterSeparateIPs(t *testing.T) {
	l := NewLoginLimiter()
	for i := 0; i <= loginMaxAttempts; i++ {
		l.Record("a")
	}
	if blocked, _ := l.Blocked("b"); blocked {
		t.Fatal("unrelated IP blocked")
	}
}

// A fresh limiter starts with no attempts, so a restart of the process
// forgives an address that was being refused.
func TestLoginLimiterStartsClean(t *testing.T) {
	l := NewLoginLimiter()
	for range loginMaxAttempts + 1 {
		l.Record("ip")
	}
	if blocked, _ := l.Blocked("ip"); !blocked {
		t.Fatal("the limiter did not refuse after the budget was spent")
	}
	fresh := NewLoginLimiter()
	if blocked, _ := fresh.Blocked("ip"); blocked {
		t.Fatal("a fresh limiter started with attempts recorded")
	}
}

// The sweep is what bounds the maps, so it is driven directly with an
// injected clock rather than by waiting for the real one.
func TestWindowSweepBoundsTheKeySet(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	w := newWindow(2, time.Minute, 10, time.Second, func() time.Time { return now })
	for i := range 50 {
		w.hit(fmt.Sprintf("key-%d", i))
		now = now.Add(time.Millisecond)
	}
	if got := w.order.Len(); got > 10 {
		t.Fatalf("tracked %d keys, want at most 10", got)
	}
	// Once the window has passed, a sweep clears what is left.
	now = now.Add(2 * time.Minute)
	w.hit("fresh")
	w.sweepLocked(now)
	if got := w.order.Len(); got != 1 {
		t.Fatalf("tracked %d keys after the sweep, want 1", got)
	}
}

// A key that keeps being refused must not grow its slice without bound.
func TestWindowCapsAHammeredKey(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	w := newWindow(2, time.Hour, 100, time.Minute, func() time.Time { return now })
	for range 5000 {
		w.hit("hammer")
		now = now.Add(time.Millisecond)
	}
	w.sweepLocked(now)
	if got := len(w.hits["hammer"]); got > 2+attemptsPerKeyLimit+1 {
		t.Fatalf("a hammered key kept %d hits", got)
	}
}