package test

import (
	"RenewCMS/api/middleware"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func newConnectionKeyLimiter(trustedProxies []string, maxAttempts int, window time.Duration) http.Handler {
	limited := middleware.NewRateLimiter(maxAttempts, window)(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	return middleware.TrustedProxyRealIP(trustedProxies)(limited)
}

func sendFromConnection(handler http.Handler, remoteAddr string, realIP string) int {
	req := httptest.NewRequest(http.MethodPost, "/v1/auth/login", nil)
	req.RemoteAddr = remoteAddr
	if realIP != "" {
		req.Header.Set("X-Real-IP", realIP)
	}
	rec := httptest.NewRecorder()
	handler.ServeHTTP(rec, req)
	return rec.Code
}

func TestRateLimitKeysOnTheConnectionHost(t *testing.T) {
	t.Run("New connections from the same address share one counter whatever their source port", func(t *testing.T) {
		handler := newConnectionKeyLimiter(nil, 5, time.Hour)

		blocked := 0
		for port := 40000; port < 40100; port++ {
			if sendFromConnection(handler, fmt.Sprintf("203.0.113.7:%d", port), "") == http.StatusTooManyRequests {
				blocked++
			}
		}

		if blocked != 95 {
			t.Errorf("expected 95 of 100 attempts blocked, got %d", blocked)
		}
	})

	t.Run("IPv6 connections from the same address share one counter", func(t *testing.T) {
		handler := newConnectionKeyLimiter(nil, 1, time.Hour)

		first := sendFromConnection(handler, "[2001:db8::1]:50000", "")
		second := sendFromConnection(handler, "[2001:db8::1]:50001", "")

		if first != http.StatusOK || second != http.StatusTooManyRequests {
			t.Errorf("expected 200 then 429, got %d then %d", first, second)
		}
	})

	t.Run("Different addresses keep independent counters", func(t *testing.T) {
		handler := newConnectionKeyLimiter(nil, 1, time.Hour)

		first := sendFromConnection(handler, "203.0.113.7:1", "")
		second := sendFromConnection(handler, "203.0.113.8:1", "")

		if first != http.StatusOK || second != http.StatusOK {
			t.Errorf("expected both 200, got %d and %d", first, second)
		}
	})

	t.Run("A forwarded client keeps one counter across the proxy's connections", func(t *testing.T) {
		handler := newConnectionKeyLimiter([]string{"10.0.0.1"}, 1, time.Hour)

		first := sendFromConnection(handler, "10.0.0.1:51000", "198.51.100.9")
		second := sendFromConnection(handler, "10.0.0.1:51001", "198.51.100.9")

		if first != http.StatusOK || second != http.StatusTooManyRequests {
			t.Errorf("expected 200 then 429, got %d then %d", first, second)
		}
	})

	t.Run("A counter starts over once its window has passed", func(t *testing.T) {
		handler := newConnectionKeyLimiter(nil, 1, 20*time.Millisecond)

		first := sendFromConnection(handler, "203.0.113.7:1", "")
		second := sendFromConnection(handler, "203.0.113.7:2", "")
		time.Sleep(30 * time.Millisecond)
		third := sendFromConnection(handler, "203.0.113.7:3", "")

		if first != http.StatusOK || second != http.StatusTooManyRequests || third != http.StatusOK {
			t.Errorf("expected 200, 429, 200, got %d, %d, %d", first, second, third)
		}
	})

	t.Run("Sweeping expired counters never resets a counter whose window is still open", func(t *testing.T) {
		handler := newConnectionKeyLimiter(nil, 1, 60*time.Millisecond)

		_ = sendFromConnection(handler, "198.51.100.1:1", "")
		time.Sleep(40 * time.Millisecond)
		first := sendFromConnection(handler, "203.0.113.7:1", "")
		time.Sleep(30 * time.Millisecond)
		second := sendFromConnection(handler, "203.0.113.7:2", "")

		if first != http.StatusOK || second != http.StatusTooManyRequests {
			t.Errorf("expected 200 then 429 across a sweep, got %d then %d", first, second)
		}
	})
}