package middleware

import (
	"net/http"
	"testing"

	"github.com/gin-gonic/gin"

	"github.com/wangjia/jiu/backend/config"
)

func newQuotaRouter(scope string, limit int) *gin.Engine {
	r := gin.New()
	r.GET("/", DailyQuota(scope, func() int { return limit }), func(c *gin.Context) {
		c.String(http.StatusOK, "ok")
	})
	return r
}

func TestDailyQuotaExceeded429(t *testing.T) {
	enableRateLimit(t)
	r := newQuotaRouter("tq1", 3)
	for i := 0; i < 3; i++ {
		if w := hit(r, "7.7.7.7"); w.Code != http.StatusOK {
			t.Fatalf("配额内第 %d 发应放行，得到 %d", i+1, w.Code)
		}
	}
	w := hit(r, "7.7.7.7")
	if w.Code != http.StatusTooManyRequests {
		t.Fatalf("超日配额应 429，得到 %d", w.Code)
	}
	if w.Header().Get("Retry-After") == "" {
		t.Fatal("429 应带 Retry-After（到午夜秒数）")
	}
	// 不同 IP 独立
	if w := hit(r, "8.8.8.8"); w.Code != http.StatusOK {
		t.Fatalf("其他 IP 应不受影响，得到 %d", w.Code)
	}
}

func TestDailyQuotaScopesIndependent(t *testing.T) {
	enableRateLimit(t)
	a := newQuotaRouter("tq2a", 1)
	b := newQuotaRouter("tq2b", 1)
	if w := hit(a, "9.9.9.9"); w.Code != http.StatusOK {
		t.Fatalf("scope a 首发应放行，得到 %d", w.Code)
	}
	if w := hit(b, "9.9.9.9"); w.Code != http.StatusOK {
		t.Fatalf("scope b 有独立配额池，应放行，得到 %d", w.Code)
	}
	if w := hit(a, "9.9.9.9"); w.Code != http.StatusTooManyRequests {
		t.Fatalf("scope a 第二发应 429，得到 %d", w.Code)
	}
}

func TestDailyQuotaDisabledOrZeroPassthrough(t *testing.T) {
	// limit=0 → 该闸关闭
	enableRateLimit(t)
	r := newQuotaRouter("tq3", 0)
	for i := 0; i < 5; i++ {
		if w := hit(r, "10.0.0.1"); w.Code != http.StatusOK {
			t.Fatalf("limit=0 第 %d 发应放行，得到 %d", i+1, w.Code)
		}
	}
	// 总开关关闭 → 放行
	prev := config.C.RateLimit.Enabled
	config.C.RateLimit.Enabled = false
	defer func() { config.C.RateLimit.Enabled = prev }()
	r2 := newQuotaRouter("tq4", 1)
	for i := 0; i < 5; i++ {
		if w := hit(r2, "10.0.0.2"); w.Code != http.StatusOK {
			t.Fatalf("总开关关闭第 %d 发应放行，得到 %d", i+1, w.Code)
		}
	}
}