package mcp_test

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

	"github.com/stretchr/testify/assert"

	"github.com/yan1h/agent-coding-workflow/internal/mcp"
)

func TestRateLimiter_AllowUntilCap(t *testing.T) {
	t.Parallel()
	now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
	limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
		PerTokenPerMinute: 3,
		GlobalPerMinute:   100,
	}, func() time.Time { return now })

	for i := 0; i < 3; i++ {
		ok, _ := limiter.Allow("token-A")
		assert.True(t, ok, "request %d should pass", i)
	}
	ok, reason := limiter.Allow("token-A")
	assert.False(t, ok)
	assert.Equal(t, "per_token", reason)
}

func TestRateLimiter_GlobalCap(t *testing.T) {
	t.Parallel()
	now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
	limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
		PerTokenPerMinute: 100,
		GlobalPerMinute:   2,
	}, func() time.Time { return now })

	ok, _ := limiter.Allow("token-A")
	assert.True(t, ok)
	ok, _ = limiter.Allow("token-B")
	assert.True(t, ok)
	ok, reason := limiter.Allow("token-C")
	assert.False(t, ok)
	assert.Equal(t, "global", reason)
}

func TestRateLimiter_WindowReset(t *testing.T) {
	t.Parallel()
	current := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
	limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
		PerTokenPerMinute: 1,
		GlobalPerMinute:   100,
	}, func() time.Time { return current })

	ok, _ := limiter.Allow("A")
	assert.True(t, ok)
	ok, _ = limiter.Allow("A")
	assert.False(t, ok)

	// 跨分钟窗口
	current = current.Add(time.Minute)
	ok, _ = limiter.Allow("A")
	assert.True(t, ok, "after window reset should allow again")
}

func TestRateLimitMiddleware_429(t *testing.T) {
	t.Parallel()
	now := time.Date(2026, 5, 7, 12, 0, 0, 0, time.UTC)
	limiter := mcp.NewRateLimiter(mcp.RateLimitConfig{
		PerTokenPerMinute: 1,
		GlobalPerMinute:   100,
	}, func() time.Time { return now })

	// 内层 handler 计数被调次数
	var hits atomic.Int32
	innerH := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		w.WriteHeader(http.StatusOK)
	})

	// 模拟 auth 中间件已注入 ctx：直接 context.WithValue + mcp.AuthContextKey
	withSession := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := context.WithValue(r.Context(), mcp.AuthContextKey, &mcp.AuthSession{
			TokenID: "fake-token-id",
			UserID:  "fake-user",
		})
		mcp.NewRateLimitMiddleware(limiter)(innerH).ServeHTTP(w, r.WithContext(ctx))
	})

	req := httptest.NewRequest("GET", "/mcp", nil)
	rec := httptest.NewRecorder()
	withSession.ServeHTTP(rec, req)
	assert.Equal(t, http.StatusOK, rec.Code)

	rec2 := httptest.NewRecorder()
	withSession.ServeHTTP(rec2, req)
	assert.Equal(t, http.StatusTooManyRequests, rec2.Code)
	assert.Equal(t, int32(1), hits.Load())
}