package ratelimit

import (
	"fmt"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

	"github.com/labstack/echo/v4"
	"golang.org/x/time/rate"
)

func TestDualLayerStore_GlobalLimit(t *testing.T) {
	config := &Config{
		GlobalRate:        rate.Limit(2),  // 2 req/s global
		GlobalBurst:       3,              // 3 burst global
		DefaultRate:       rate.Limit(10), // 10 req/s per IP (higher than global)
		DefaultBurst:      20,             // 20 burst per IP
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	store := NewDualLayerStore(config)

	// Should allow first 3 requests (burst)
	for i := 0; i < 3; i++ {
		allowed, err := store.Allow("192.168.1.1")
		if err != nil || !allowed {
			t.Fatalf("Request %d should be allowed (burst): err=%v, allowed=%v", i+1, err, allowed)
		}
	}

	// 4th request should exceed global limit
	allowed, err := store.Allow("192.168.1.1")
	if allowed || err == nil {
		t.Error("4th request should exceed global rate limit")
	}
}

func TestDualLayerStore_PerIPLimit(t *testing.T) {
	config := &Config{
		GlobalRate:        rate.Limit(1000), // High global limit (won't hit it)
		GlobalBurst:       2000,
		DefaultRate:       rate.Limit(5), // 5 req/s per IP
		DefaultBurst:      3,             // 3 burst per IP
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	store := NewDualLayerStore(config)

	// Should allow first 3 requests (burst)
	for i := 0; i < 3; i++ {
		allowed, err := store.Allow("192.168.1.1")
		if err != nil || !allowed {
			t.Fatalf("Request %d should be allowed (burst): err=%v, allowed=%v", i+1, err, allowed)
		}
	}

	// 4th request should exceed per-IP limit
	allowed, err := store.Allow("192.168.1.1")
	if allowed {
		t.Error("4th request should exceed per-IP rate limit")
	}
	if err != nil {
		t.Logf("Per-IP limit error (expected): %v", err)
	}

	// Different IP should still be allowed
	allowed, err = store.Allow("192.168.1.2")
	if err != nil || !allowed {
		t.Error("Request from different IP should be allowed")
	}
}

func TestDualLayerStore_MultipleIPs(t *testing.T) {
	config := &Config{
		GlobalRate:        rate.Limit(1000),
		GlobalBurst:       2000,
		DefaultRate:       rate.Limit(5),
		DefaultBurst:      2,
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	store := NewDualLayerStore(config)

	// Each IP should have independent limits
	ips := []string{"192.168.1.1", "192.168.1.2", "192.168.1.3"}

	for _, ip := range ips {
		// Each IP can make burst requests
		for i := 0; i < 2; i++ {
			allowed, err := store.Allow(ip)
			if err != nil || !allowed {
				t.Errorf("IP %s request %d should be allowed: err=%v, allowed=%v", ip, i+1, err, allowed)
			}
		}

		// 3rd request should exceed per-IP limit
		allowed, err := store.Allow(ip)
		if allowed {
			t.Errorf("IP %s request 3 should exceed per-IP limit", ip)
		}
		if err != nil {
			t.Logf("Per-IP limit for %s (expected): %v", ip, err)
		}
	}
}

func TestNewRateLimitMiddleware_DefaultLimit(t *testing.T) {
	config := &Config{
		GlobalRate:        rate.Limit(1000),
		GlobalBurst:       2000,
		DefaultRate:       rate.Limit(5),
		DefaultBurst:      3,
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	middleware := NewRateLimitMiddleware(config, logger)

	e := echo.New()
	e.Use(middleware)

	// Test handler
	e.GET("/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	})

	// Should allow first 3 requests (burst)
	for i := 0; i < 3; i++ {
		req := httptest.NewRequest(http.MethodGet, "/test", nil)
		req.Header.Set("X-Real-IP", "192.168.1.1")
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, req)

		if rec.Code != http.StatusOK {
			t.Fatalf("Request %d should succeed: got status %d", i+1, rec.Code)
		}

		// Check RateLimit header
		rateLimitHeader := rec.Header().Get("RateLimit")
		if rateLimitHeader == "" {
			t.Errorf("Request %d should have RateLimit header", i+1)
		}
	}

	// 4th request should be rate limited
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Header.Set("X-Real-IP", "192.168.1.1")
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	if rec.Code != http.StatusTooManyRequests {
		t.Errorf("4th request should be rate limited: got status %d", rec.Code)
	}

	// Check Retry-After header
	retryAfter := rec.Header().Get("Retry-After")
	if retryAfter == "" {
		t.Error("429 response should have Retry-After header")
	}

	// Check RateLimit header
	rateLimit := rec.Header().Get("RateLimit")
	if rateLimit == "" {
		t.Error("429 response should have RateLimit header")
	}
}

func TestNewRateLimitMiddleware_EndpointOverride(t *testing.T) {
	config := &Config{
		GlobalRate:   rate.Limit(1000),
		GlobalBurst:  2000,
		DefaultRate:  rate.Limit(5), // Lower default to test override headers
		DefaultBurst: 3,
		ExpiresIn:    3 * time.Minute,
		EndpointOverrides: map[string]EndpointLimit{
			"POST /restricted": {
				GlobalRate:  rate.Limit(100),
				GlobalBurst: 200,
				Rate:        rate.Limit(1),
				Burst:       2,
			},
		},
	}

	logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
	middleware := NewRateLimitMiddleware(config, logger)

	e := echo.New()
	e.Use(middleware)

	e.POST("/restricted", func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	})

	// Note: Current implementation uses default limits for actual rate limiting
	// Endpoint overrides only affect headers (RateLimit, Retry-After)
	// This is a known limitation that could be enhanced in the future

	// Should allow first 3 requests (default burst)
	for i := 0; i < 3; i++ {
		req := httptest.NewRequest(http.MethodPost, "/restricted", nil)
		req.Header.Set("X-Real-IP", "192.168.1.1")
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, req)

		if rec.Code != http.StatusOK {
			t.Fatalf("Request %d should succeed: got status %d", i+1, rec.Code)
		}

		// Verify endpoint override affects the RateLimit header
		rateLimitHeader := rec.Header().Get("RateLimit")
		// Override has Burst=2, Rate=1, so window=2/1=2
		expectedHeader := "2;window=2"
		if rateLimitHeader != expectedHeader {
			t.Errorf("Request %d: Expected RateLimit header %s, got %s", i+1, expectedHeader, rateLimitHeader)
		}
	}

	// 4th request should be rate limited (default burst=3)
	req := httptest.NewRequest(http.MethodPost, "/restricted", nil)
	req.Header.Set("X-Real-IP", "192.168.1.1")
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	if rec.Code != http.StatusTooManyRequests {
		t.Errorf("4th request should be rate limited: got status %d", rec.Code)
	}
}

func TestNewRateLimitMiddleware_HeaderFormat(t *testing.T) {
	config := &Config{
		GlobalRate:        rate.Limit(1000),
		GlobalBurst:       2000,
		DefaultRate:       rate.Limit(10),
		DefaultBurst:      20,
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	middleware := NewRateLimitMiddleware(config, logger)

	e := echo.New()
	e.Use(middleware)

	e.GET("/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	})

	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Header.Set("X-Real-IP", "192.168.1.1")
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	// Check RateLimit header format
	rateLimit := rec.Header().Get("RateLimit")
	if rateLimit == "" {
		t.Fatal("Response should have RateLimit header")
	}

	// Should match format: "{burst};window={seconds}"
	// For DefaultRate=10, DefaultBurst=20: window = 20/10 = 2
	expected := "20;window=2"
	if rateLimit != expected {
		t.Errorf("RateLimit header format incorrect: expected=%s, got=%s", expected, rateLimit)
	}
}

func TestNewRateLimitMiddleware_DifferentIPs(t *testing.T) {
	config := &Config{
		GlobalRate:        rate.Limit(1000),
		GlobalBurst:       2000,
		DefaultRate:       rate.Limit(5),
		DefaultBurst:      2,
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
	middleware := NewRateLimitMiddleware(config, logger)

	e := echo.New()
	e.Use(middleware)

	e.GET("/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	})

	// IP1 exhausts its limit
	for i := 0; i < 2; i++ {
		req := httptest.NewRequest(http.MethodGet, "/test", nil)
		req.Header.Set("X-Real-IP", "192.168.1.1")
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, req)

		if rec.Code != http.StatusOK {
			t.Fatalf("IP1 request %d should succeed", i+1)
		}
	}

	// IP1's 3rd request should be rate limited
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Header.Set("X-Real-IP", "192.168.1.1")
	rec := httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	if rec.Code != http.StatusTooManyRequests {
		t.Error("IP1's 3rd request should be rate limited")
	}

	// IP2 should still be allowed
	req = httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Header.Set("X-Real-IP", "192.168.1.2")
	rec = httptest.NewRecorder()
	e.ServeHTTP(rec, req)

	if rec.Code != http.StatusOK {
		t.Errorf("IP2's first request should succeed: got status %d", rec.Code)
	}
}

func TestSetRateLimitHeader(t *testing.T) {
	tests := []struct {
		name          string
		config        *Config
		route         string
		expectedValue string
	}{
		{
			name: "default config",
			config: &Config{
				DefaultRate:       rate.Limit(10),
				DefaultBurst:      20,
				EndpointOverrides: make(map[string]EndpointLimit),
			},
			route:         "GET /test",
			expectedValue: "20;window=2",
		},
		{
			name: "endpoint override",
			config: &Config{
				DefaultRate:  rate.Limit(10),
				DefaultBurst: 20,
				EndpointOverrides: map[string]EndpointLimit{
					"POST /upload": {
						Rate:  rate.Limit(2),
						Burst: 5,
					},
				},
			},
			route:         "POST /upload",
			expectedValue: "5;window=2",
		},
		{
			name: "high rate (window=1)",
			config: &Config{
				DefaultRate:       rate.Limit(100),
				DefaultBurst:      50,
				EndpointOverrides: make(map[string]EndpointLimit),
			},
			route:         "GET /test",
			expectedValue: "50;window=1", // 50/100 = 0.5, rounded to 1
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := echo.New()
			req := httptest.NewRequest(http.MethodGet, "/", nil)
			rec := httptest.NewRecorder()
			c := e.NewContext(req, rec)

			setRateLimitHeader(c, tt.config, tt.route)

			value := rec.Header().Get("RateLimit")
			if value != tt.expectedValue {
				t.Errorf("Expected RateLimit=%s, got %s", tt.expectedValue, value)
			}
		})
	}
}

func TestDualLayerStore_Cleanup(t *testing.T) {
	config := &Config{
		GlobalRate:        rate.Limit(1000),
		GlobalBurst:       2000,
		DefaultRate:       rate.Limit(10),
		DefaultBurst:      20,
		ExpiresIn:         100 * time.Millisecond, // Short expiry for testing
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	store := NewDualLayerStore(config)

	// Mock time function
	currentTime := time.Now()
	store.timeNow = func() time.Time {
		return currentTime
	}

	// Add some IPs
	_, _ = store.Allow("192.168.1.1")
	_, _ = store.Allow("192.168.1.2")
	_, _ = store.Allow("192.168.1.3")

	initialCount := len(store.ipLimiters)
	if initialCount != 3 {
		t.Errorf("Expected 3 IP limiters, got %d", initialCount)
	}

	// Advance time past expiry
	currentTime = currentTime.Add(200 * time.Millisecond)

	// Trigger cleanup by making another request
	_, _ = store.Allow("192.168.1.4")

	// Note: Current cleanup implementation clears all limiters (simplified)
	// In production, this would track last access time per IP
	// For now, we just verify cleanup was triggered (limiters were reset)
	finalCount := len(store.ipLimiters)
	if finalCount > initialCount {
		t.Errorf("Expected cleanup to have occurred, but limiter count increased from %d to %d", initialCount, finalCount)
	}
}

func BenchmarkDualLayerStore_Allow(b *testing.B) {
	config := &Config{
		GlobalRate:        rate.Limit(1000),
		GlobalBurst:       2000,
		DefaultRate:       rate.Limit(100),
		DefaultBurst:      200,
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	store := NewDualLayerStore(config)
	ip := "192.168.1.1"

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = store.Allow(ip)
	}
}

func BenchmarkMiddleware_Request(b *testing.B) {
	config := &Config{
		GlobalRate:        rate.Limit(10000), // High limit for benchmarking
		GlobalBurst:       20000,
		DefaultRate:       rate.Limit(1000),
		DefaultBurst:      2000,
		ExpiresIn:         3 * time.Minute,
		EndpointOverrides: make(map[string]EndpointLimit),
	}

	logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
	middleware := NewRateLimitMiddleware(config, logger)

	e := echo.New()
	e.Use(middleware)
	e.GET("/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	})

	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Header.Set("X-Real-IP", "192.168.1.1")

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rec := httptest.NewRecorder()
		e.ServeHTTP(rec, req)
	}
}

func ExampleNewRateLimitMiddleware() {
	// Create config
	config := &Config{
		GlobalRate:   rate.Limit(1000),
		GlobalBurst:  2000,
		DefaultRate:  rate.Limit(10),
		DefaultBurst: 20,
		ExpiresIn:    3 * time.Minute,
		EndpointOverrides: map[string]EndpointLimit{
			"POST /documents": {
				GlobalRate:  rate.Limit(50),
				GlobalBurst: 100,
				Rate:        rate.Limit(5),
				Burst:       10,
			},
		},
	}

	// Create logger
	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))

	// Create middleware
	middleware := NewRateLimitMiddleware(config, logger)

	// Use with Echo
	e := echo.New()
	e.Use(middleware)

	e.GET("/api/test", func(c echo.Context) error {
		return c.String(http.StatusOK, "success")
	})

	fmt.Println("Rate limiting enabled")
	// Output: Rate limiting enabled
}