package middleware

import (
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
)

// newRateLimitRouter 构造一个使用 RateLimit 中间件、所有请求都返回 200 的测试 router。
// 返回 router 与下游被调用次数计数器，便于断言请求是否被放行。
func newRateLimitRouter(rps float64, burst int) (*gin.Engine, *int32) {
	gin.SetMode(gin.TestMode)
	var count int32
	r := gin.New()
	r.Use(RateLimit(rps, burst))
	r.GET("/test", func(c *gin.Context) {
		atomic.AddInt32(&count, 1)
		c.Status(http.StatusOK)
	})
	return r, &count
}

// doRequest 发起一次测试请求，返回响应状态码。
func doRequest(r *gin.Engine, remoteAddr string) int {
	w := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.RemoteAddr = remoteAddr
	r.ServeHTTP(w, req)
	return w.Code
}

// TestRateLimit_AllowUnderBurst 验证突发请求 ≤ burst 时全部放行。
func TestRateLimit_AllowUnderBurst(t *testing.T) {
	r, count := newRateLimitRouter(0.1 /*rps*/, 5 /*burst*/)

	// 5 个请求应全部放行（初始满桶 5 令牌）
	for i := 0; i < 5; i++ {
		if code := doRequest(r, "1.1.1.1:1234"); code != http.StatusOK {
			t.Fatalf("请求 %d 应放行（200），实际 %d", i+1, code)
		}
	}
	if got := atomic.LoadInt32(count); got != 5 {
		t.Fatalf("下游应被调用 5 次，实际 %d", got)
	}
}

// TestRateLimit_RejectOverBurst 验证突发请求 > burst 时返回 429。
func TestRateLimit_RejectOverBurst(t *testing.T) {
	r, count := newRateLimitRouter(0.1 /*rps*/, 3 /*burst*/)

	// 前 3 个放行
	for i := 0; i < 3; i++ {
		if code := doRequest(r, "2.2.2.2:1234"); code != http.StatusOK {
			t.Fatalf("请求 %d 应放行（200），实际 %d", i+1, code)
		}
	}
	// 第 4 个应被限流（429）
	if code := doRequest(r, "2.2.2.2:1234"); code != http.StatusTooManyRequests {
		t.Fatalf("第 4 个请求应被限流（429），实际 %d", code)
	}
	if got := atomic.LoadInt32(count); got != 3 {
		t.Fatalf("下游应被调用 3 次，实际 %d", got)
	}
}

// TestRateLimit_RefillTokens 验证令牌按 rps 速率补充。
// rps=10 表示每秒补 10 个令牌，即每 100ms 补 1 个；消耗 1 个后等待 ~150ms 应补回 1+ 个。
func TestRateLimit_RefillTokens(t *testing.T) {
	r, count := newRateLimitRouter(10 /*rps*/, 1 /*burst*/)

	// 第 1 个放行（消耗唯一令牌）
	if code := doRequest(r, "3.3.3.3:1234"); code != http.StatusOK {
		t.Fatalf("第 1 个请求应放行（200），实际 %d", code)
	}
	// 立即第 2 个应被限流（无令牌）
	if code := doRequest(r, "3.3.3.3:1234"); code != http.StatusTooManyRequests {
		t.Fatalf("第 2 个请求应被限流（429），实际 %d", code)
	}
	// 等待 150ms，应补充 1.5 个令牌（capped 至 burst=1）
	time.Sleep(150 * time.Millisecond)
	// 第 3 个应再次放行
	if code := doRequest(r, "3.3.3.3:1234"); code != http.StatusOK {
		t.Fatalf("等待补充后请求应放行（200），实际 %d", code)
	}
	if got := atomic.LoadInt32(count); got != 2 {
		t.Fatalf("下游应被调用 2 次，实际 %d", got)
	}
}

// TestRateLimit_PerIPIsolation 验证不同 IP 的桶相互独立。
func TestRateLimit_PerIPIsolation(t *testing.T) {
	r, count := newRateLimitRouter(0.1 /*rps*/, 2 /*burst*/)

	// IP A 消耗 2 个令牌（满桶）
	for i := 0; i < 2; i++ {
		if code := doRequest(r, "10.0.0.1:1234"); code != http.StatusOK {
			t.Fatalf("IP A 请求 %d 应放行，实际 %d", i+1, code)
		}
	}
	// IP A 第 3 个应被限流
	if code := doRequest(r, "10.0.0.1:1234"); code != http.StatusTooManyRequests {
		t.Fatalf("IP A 第 3 个请求应被限流，实际 %d", code)
	}
	// IP B 仍有独立满桶，应放行
	if code := doRequest(r, "10.0.0.2:1234"); code != http.StatusOK {
		t.Fatalf("IP B 第 1 个请求应放行（独立桶），实际 %d", code)
	}
	if got := atomic.LoadInt32(count); got != 3 {
		t.Fatalf("下游应被调用 3 次（A×2 + B×1），实际 %d", got)
	}
}

// TestRateLimit_CleanupExpiredBuckets 验证 10 分钟无访问的桶被清理。
// 使用短周期参数（50ms 间隔、500ms 过期）加速测试；直接构造 rateLimiter 而不经过 RateLimit
// （后者会启动 5min/10min 的清理 goroutine，无法在测试中观察）。
// 时序：创建旧桶（1h 前）→ 启动 cleanup → 等 100ms（旧桶被清，无新桶）→ 创建新桶 → 等 100ms（新桶未过 500ms，保留）。
func TestRateLimit_CleanupExpiredBuckets(t *testing.T) {
	rl := &rateLimiter{rps: 1, burst: 1}

	// 创建旧桶并人为把 lastTime 调到 1 小时前
	b := rl.getBucket("9.9.9.9:1234")
	b.mu.Lock()
	b.lastTime = time.Now().Add(-1 * time.Hour)
	b.mu.Unlock()

	// 启动 cleanup，50ms 间隔、500ms 过期
	go rl.cleanup(50*time.Millisecond, 500*time.Millisecond)

	// 等 100ms 让 cleanup 执行（旧桶 lastTime=1h 前，远超 500ms，必被清）
	time.Sleep(100 * time.Millisecond)

	// 旧桶应已清理
	if _, ok := rl.buckets.Load("9.9.9.9:1234"); ok {
		t.Fatal("1 小时无访问的旧桶应被清理")
	}

	// 此时创建新鲜桶（lastTime = now）
	rl.getBucket("8.8.8.8:1234")

	// 再等 100ms 让 cleanup 再执行
	time.Sleep(100 * time.Millisecond)

	// 新鲜桶 lastTime 仅 ~100ms 前，未超 500ms maxIdle，应保留
	if _, ok := rl.buckets.Load("8.8.8.8:1234"); !ok {
		t.Fatal("新鲜桶（<500ms）不应被清理")
	}
}