package billing

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

	"github.com/gin-gonic/gin"
	"github.com/go-logr/logr"
	"github.com/tidbcloud/devops/pkg/infra/config"
	"github.com/tidbcloud/devops/pkg/infra/consts"
	"golang.org/x/time/rate"
)

func TestProxyRateLimiter_Disabled(t *testing.T) {
	gin.SetMode(gin.TestMode)
	lim := newProxyRateLimiter(config.BillingProxyRateLimitConfig{Enabled: false}, logr.Discard())
	r := gin.New()
	r.GET("/billing-proxy/*any",
		setOperator(consts.CtxOperatorTypeAPIToken, "agent@example.com"),
		lim.Middleware(),
		func(c *gin.Context) { c.Status(http.StatusOK) },
	)

	for i := 0; i < 5; i++ {
		w := httptest.NewRecorder()
		req, _ := http.NewRequest(http.MethodGet, "/billing-proxy/foo", nil)
		r.ServeHTTP(w, req)
		if w.Code != http.StatusOK {
			t.Fatalf("request %d: got status %d, want %d", i, w.Code, http.StatusOK)
		}
	}
}

func TestProxyRateLimiter_EnforcedPerOperator(t *testing.T) {
	gin.SetMode(gin.TestMode)
	lim := newProxyRateLimiter(config.BillingProxyRateLimitConfig{
		Enabled: true,
		APIToken: config.RateLimitBucketConfig{
			RPS:   100,
			Burst: 2,
		},
	}, logr.Discard())
	r := gin.New()
	r.GET("/billing-proxy/*any",
		setOperator(consts.CtxOperatorTypeAPIToken, "agent@example.com"),
		lim.Middleware(),
		func(c *gin.Context) { c.Status(http.StatusOK) },
	)

	for i := 0; i < 2; i++ {
		w := httptest.NewRecorder()
		req, _ := http.NewRequest(http.MethodGet, "/billing-proxy/foo", nil)
		r.ServeHTTP(w, req)
		if w.Code != http.StatusOK {
			t.Fatalf("request %d: got status %d, want %d", i, w.Code, http.StatusOK)
		}
	}

	w := httptest.NewRecorder()
	req, _ := http.NewRequest(http.MethodGet, "/billing-proxy/foo", nil)
	r.ServeHTTP(w, req)
	if w.Code != http.StatusTooManyRequests {
		t.Fatalf("got status %d, want %d", w.Code, http.StatusTooManyRequests)
	}
}

func TestProxyRateLimiter_SeparateOperators(t *testing.T) {
	gin.SetMode(gin.TestMode)
	lim := newProxyRateLimiter(config.BillingProxyRateLimitConfig{
		Enabled: true,
		APIToken: config.RateLimitBucketConfig{
			RPS:   100,
			Burst: 1,
		},
	}, logr.Discard())
	r := gin.New()
	r.GET("/billing-proxy/*any",
		func(c *gin.Context) {
			if c.GetHeader("X-Operator") != "" {
				setOperator(consts.CtxOperatorTypeAPIToken, c.GetHeader("X-Operator"))(c)
			}
		},
		lim.Middleware(),
		func(c *gin.Context) { c.Status(http.StatusOK) },
	)

	for _, op := range []string{"a@example.com", "b@example.com"} {
		w := httptest.NewRecorder()
		req, _ := http.NewRequest(http.MethodGet, "/billing-proxy/foo", nil)
		req.Header.Set("X-Operator", op)
		r.ServeHTTP(w, req)
		if w.Code != http.StatusOK {
			t.Fatalf("operator %s: got status %d, want %d", op, w.Code, http.StatusOK)
		}
	}
}

func TestProxyRateLimiter_ClearMapWhenExceedsCapacity(t *testing.T) {
	lim := newProxyRateLimiter(config.BillingProxyRateLimitConfig{Enabled: true}, logr.Discard())

	lim.mu.Lock()
	for i := 0; i < maxLimiterEntries; i++ {
		key := fmt.Sprintf("%s:user%d@example.com", consts.CtxOperatorTypeAPIToken, i)
		lim.entries[key] = rate.NewLimiter(rate.Limit(1), 1)
	}
	lim.mu.Unlock()

	lim.getLimiter("api_token:new@example.com", consts.CtxOperatorTypeAPIToken)

	lim.mu.Lock()
	defer lim.mu.Unlock()
	if len(lim.entries) != 1 {
		t.Fatalf("got %d entries after eviction, want 1", len(lim.entries))
	}
	if _, ok := lim.entries["api_token:new@example.com"]; !ok {
		t.Fatal("expected new operator limiter to remain after eviction")
	}
}

func TestLimitsForOperatorType(t *testing.T) {
	cfg := config.BillingProxyRateLimitConfig{
		APIToken: config.RateLimitBucketConfig{RPS: 1, Burst: 5},
		User:     config.RateLimitBucketConfig{RPS: 3, Burst: 9},
	}

	rps, burst := limitsForOperatorType(cfg, consts.CtxOperatorTypeAPIToken)
	if rps != 1 || burst != 5 {
		t.Fatalf("api_token: got rps=%v burst=%d, want 1/5", rps, burst)
	}

	rps, burst = limitsForOperatorType(cfg, consts.CtxOperatorTypeUser)
	if rps != 3 || burst != 9 {
		t.Fatalf("user: got rps=%v burst=%d, want 3/9", rps, burst)
	}

	// Unknown operator types use the user bucket; zero config falls back to code defaults.
	rps, burst = limitsForOperatorType(config.BillingProxyRateLimitConfig{}, "unknown")
	if rps != 2 || burst != 30 {
		t.Fatalf("default user: got rps=%v burst=%d, want 2/30", rps, burst)
	}
}

func setOperator(operatorType, operator string) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Set(consts.CtxOperatorType, operatorType)
		c.Set(consts.CtxOperator, operator)
	}
}