package middleware

import (
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestRateLimiter_Allow(t *testing.T) {
	rl := newRateLimiter(3)

	for i := 0; i < 3; i++ {
		if !rl.allow("1.2.3.4") {
			t.Errorf("zahtev %d bi trebalo da prodje", i+1)
		}
	}

	if rl.allow("1.2.3.4") {
		t.Error("4. zahtev ne sme proci (limit 3)")
	}
}

func TestRateLimiter_DifferentIPs(t *testing.T) {
	rl := newRateLimiter(2)

	rl.allow("1.1.1.1")
	rl.allow("1.1.1.1")

	// Razlicit IP — mora proci
	if !rl.allow("2.2.2.2") {
		t.Error("razlicit IP mora proci")
	}
}

func TestRateLimit_Middleware(t *testing.T) {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	limited := RateLimit(2)(handler)

	for i := 0; i < 2; i++ {
		req := httptest.NewRequest("POST", "/test", nil)
		req.RemoteAddr = "10.0.0.1:1234"
		w := httptest.NewRecorder()
		limited.ServeHTTP(w, req)

		if w.Code != http.StatusOK {
			t.Errorf("zahtev %d: ocekivan 200, dobijen %d", i+1, w.Code)
		}
	}

	// 3rd request should be rate limited
	req := httptest.NewRequest("POST", "/test", nil)
	req.RemoteAddr = "10.0.0.1:1234"
	w := httptest.NewRecorder()
	limited.ServeHTTP(w, req)

	if w.Code != http.StatusTooManyRequests {
		t.Errorf("3. zahtev: ocekivan 429, dobijen %d", w.Code)
	}
}

func TestRateLimit_XForwardedFor(t *testing.T) {
	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	limited := RateLimit(1)(handler)

	req := httptest.NewRequest("POST", "/test", nil)
	req.RemoteAddr = "10.0.0.1:1234"
	req.Header.Set("X-Forwarded-For", "5.5.5.5")
	w := httptest.NewRecorder()
	limited.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Errorf("prvi zahtev mora proci")
	}

	// Second from same X-Forwarded-For
	req2 := httptest.NewRequest("POST", "/test", nil)
	req2.RemoteAddr = "10.0.0.2:5678" // razlicit RemoteAddr
	req2.Header.Set("X-Forwarded-For", "5.5.5.5") // isti X-FF
	w2 := httptest.NewRecorder()
	limited.ServeHTTP(w2, req2)

	if w2.Code != http.StatusTooManyRequests {
		t.Errorf("drugi zahtev sa istim XFF mora biti blokiran, dobijen %d", w2.Code)
	}
}