package middleware

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// fakeQuotaStore 内存实现 quotaStore, 记录每个 key 的计数与 TTL。
type fakeQuotaStore struct {
	mu      sync.Mutex
	counts  map[string]int64
	ttls    map[string]time.Duration
	incrErr error
}

func newFakeQuotaStore() *fakeQuotaStore {
	return &fakeQuotaStore{
		counts: make(map[string]int64),
		ttls:   make(map[string]time.Duration),
	}
}

func (f *fakeQuotaStore) Incr(_ context.Context, key string) (int64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.incrErr != nil {
		return 0, f.incrErr
	}
	f.counts[key]++
	return f.counts[key], nil
}

func (f *fakeQuotaStore) Expire(_ context.Context, key string, ttl time.Duration) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.ttls[key] = ttl
	return nil
}

func (f *fakeQuotaStore) totalKeys() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return len(f.counts)
}

// setupQuotaRouter 装配 quota 中间件; withKeyID 模拟 AgentAuth 已注入 api_key_id。
// 返回 router 与下游 handler 命中次数(用于验证 Abort)。
func setupQuotaRouter(store quotaStore, enabled bool, rpm int, withKeyID bool) (*gin.Engine, *int) {
	gin.SetMode(gin.TestMode)
	hits := new(int)
	r := gin.New()
	r.Use(func(c *gin.Context) {
		if withKeyID {
			c.Set("api_key_id", uint64(42))
		}
		c.Next()
	})
	r.Use(quotaWithStore(store, enabled, rpm))
	r.GET("/ping", func(c *gin.Context) {
		*hits++
		c.JSON(http.StatusOK, gin.H{"ok": true})
	})
	return r, hits
}

func doQuotaRequest(t *testing.T, r *gin.Engine) (*httptest.ResponseRecorder, map[string]interface{}) {
	t.Helper()
	w := httptest.NewRecorder()
	req := httptest.NewRequest(http.MethodGet, "/ping", nil)
	r.ServeHTTP(w, req)
	var body map[string]interface{}
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
	return w, body
}

// disabled → 直通不计数。
func TestQuotaDisabledPassthrough(t *testing.T) {
	store := newFakeQuotaStore()
	r, hits := setupQuotaRouter(store, false, 10, true)

	for i := 0; i < 3; i++ {
		w, _ := doQuotaRequest(t, r)
		assert.Equal(t, http.StatusOK, w.Code)
	}
	assert.Equal(t, 3, *hits)
	assert.Equal(t, 0, store.totalKeys())
}

// enabled 但 rpm=0 → 直通(warn 一次)。
func TestQuotaZeroRPMPassthrough(t *testing.T) {
	store := newFakeQuotaStore()
	r, hits := setupQuotaRouter(store, true, 0, true)

	for i := 0; i < 3; i++ {
		w, _ := doQuotaRequest(t, r)
		assert.Equal(t, http.StatusOK, w.Code)
	}
	assert.Equal(t, 3, *hits)
	assert.Equal(t, 0, store.totalKeys())
}

// 无 api_key_id(人类 JWT) → 直通。
func TestQuotaNoAPIKeyPassthrough(t *testing.T) {
	store := newFakeQuotaStore()
	r, hits := setupQuotaRouter(store, true, 1, false)

	for i := 0; i < 3; i++ {
		w, _ := doQuotaRequest(t, r)
		assert.Equal(t, http.StatusOK, w.Code)
	}
	assert.Equal(t, 3, *hits)
	assert.Equal(t, 0, store.totalKeys())
}

// 限额内 → 放行且计数, 首个请求设置 TTL。
func TestQuotaWithinLimit(t *testing.T) {
	store := newFakeQuotaStore()
	r, hits := setupQuotaRouter(store, true, 2, true)

	for i := 0; i < 2; i++ {
		w, _ := doQuotaRequest(t, r)
		assert.Equal(t, http.StatusOK, w.Code)
	}
	assert.Equal(t, 2, *hits)
	require.Equal(t, 1, store.totalKeys())
	for key, n := range store.counts {
		assert.Equal(t, int64(2), n)
		assert.Contains(t, key, "quota:agent:42:")
		assert.Equal(t, quotaKeyTTL, store.ttls[key])
	}
}

// 超限 → body code 429 + Retry-After: 60 + Abort(下游不执行)。
func TestQuotaExceeded(t *testing.T) {
	store := newFakeQuotaStore()
	r, hits := setupQuotaRouter(store, true, 2, true)

	for i := 0; i < 2; i++ {
		w, _ := doQuotaRequest(t, r)
		assert.Equal(t, http.StatusOK, w.Code)
	}
	w, body := doQuotaRequest(t, r)

	// 项目信封: HTTP 200, 429 体现在 body code。
	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, float64(http.StatusTooManyRequests), body["code"])
	assert.Equal(t, "rate limit exceeded", body["message"])
	assert.Equal(t, "60", w.Header().Get("Retry-After"))
	assert.Equal(t, 2, *hits, "超限请求应 Abort, 不进入下游 handler")
}

// 分钟窗口切换后重新计数。
func TestQuotaWindowReset(t *testing.T) {
	base := time.Date(2026, 7, 27, 12, 0, 30, 0, time.UTC)
	current := base
	now = func() time.Time { return current }
	defer func() { now = time.Now }()

	store := newFakeQuotaStore()
	r, hits := setupQuotaRouter(store, true, 1, true)

	// 窗口内: 第 1 次放行, 第 2 次 429。
	w, _ := doQuotaRequest(t, r)
	assert.Equal(t, http.StatusOK, w.Code)
	_, body := doQuotaRequest(t, r)
	assert.Equal(t, float64(http.StatusTooManyRequests), body["code"])
	assert.Equal(t, 1, *hits)

	// 跨入下一分钟窗口: 重新计数, 再次放行。
	current = base.Add(61 * time.Second)
	w, _ = doQuotaRequest(t, r)
	assert.Equal(t, http.StatusOK, w.Code)
	assert.Equal(t, 2, *hits)
	assert.Equal(t, 2, store.totalKeys(), "窗口切换后应使用新 key 重新计数")
}