package middleware

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

	"hfb_sys/backend/internal/modules/auth"

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

func runRateLimitedRequest(router *gin.Engine, bearerToken string) int {
	req := httptest.NewRequest(http.MethodGet, "/ping", nil)
	if bearerToken != "" {
		req.Header.Set("Authorization", "Bearer "+bearerToken)
	}
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, req)
	return recorder.Code
}

func newRateLimitRouter(limit int, rdb interface{}, handler gin.HandlerFunc) *gin.Engine {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	router.Use(handler)
	router.GET("/ping", func(c *gin.Context) {
		c.Status(http.StatusOK)
	})
	return router
}

func adminAccessToken(t *testing.T, jwtManager *auth.JWTManager, adminID uint64, phone string) string {
	t.Helper()
	pair, err := jwtManager.GenerateSubjectPairWithVersion(adminID, phone, "admin", 0)
	if err != nil {
		t.Fatalf("生成后台令牌失败: %v", err)
	}
	return pair.AccessToken
}

func TestAdminAwareRateLimitIsolatesAdminsBehindSharedIP(t *testing.T) {
	jwtManager := auth.NewJWTManager("test-secret")
	adminA := adminAccessToken(t, jwtManager, 101, "13800000001")
	adminB := adminAccessToken(t, jwtManager, 102, "13800000002")
	router := newRateLimitRouter(2, nil, AdminAwareRateLimitPerMinute(2, nil, jwtManager))

	for i := 0; i < 2; i++ {
		if code := runRateLimitedRequest(router, adminA); code != http.StatusOK {
			t.Fatalf("客服A 第 %d 次请求 = %d, want 200", i+1, code)
		}
	}
	if code := runRateLimitedRequest(router, adminA); code != http.StatusTooManyRequests {
		t.Fatalf("客服A 超额请求 = %d, want 429", code)
	}
	if code := runRateLimitedRequest(router, adminB); code != http.StatusOK {
		t.Fatalf("同出口IP的客服B 应不受客服A影响, got %d, want 200", code)
	}
}

func TestAdminAwareRateLimitFallsBackToIP(t *testing.T) {
	jwtManager := auth.NewJWTManager("test-secret")
	userPair, err := jwtManager.GenerateSubjectPairWithVersion(201, "13900000001", "user", 0)
	if err != nil {
		t.Fatalf("生成用户令牌失败: %v", err)
	}
	router := newRateLimitRouter(2, nil, AdminAwareRateLimitPerMinute(2, nil, jwtManager))

	for i := 0; i < 2; i++ {
		if code := runRateLimitedRequest(router, ""); code != http.StatusOK {
			t.Fatalf("匿名请求 第 %d 次 = %d, want 200", i+1, code)
		}
	}
	if code := runRateLimitedRequest(router, ""); code != http.StatusTooManyRequests {
		t.Fatalf("匿名请求超额 = %d, want 429", code)
	}
	// 用户令牌无法按 admin 身份解析，应与匿名共享同一个 IP 桶。
	if code := runRateLimitedRequest(router, userPair.AccessToken); code != http.StatusTooManyRequests {
		t.Fatalf("用户令牌应回退到IP维度, got %d, want 429", code)
	}
}

func TestLegacyRateLimitPerMinuteKeepsIPKey(t *testing.T) {
	router := newRateLimitRouter(1, nil, RateLimitPerMinute(1, nil))
	if code := runRateLimitedRequest(router, ""); code != http.StatusOK {
		t.Fatalf("首次匿名请求 = %d, want 200", code)
	}
	if code := runRateLimitedRequest(router, ""); code != http.StatusTooManyRequests {
		t.Fatalf("第二次匿名请求 = %d, want 429", code)
	}
}