package mcp

import (
	"net/http"
	"sync"
	"time"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	httpx "github.com/yan1h/agent-coding-workflow/internal/transport/http"
	"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
)

// RateLimiter 是一对 token bucket：per-token + 全局。
//
// 实现选择：固定时间窗口 + 计数器（每分钟重置）。简单、无第三方依赖、
// 内存占用极低。10–50 人规模 + 默认 100 req/min/token 不需要更精细的
// sliding window / leaky bucket。
//
// 并发：sync.Mutex + map；rate-limit 调用是热路径（每请求一次），但 mutex
// 持有时间极短（map lookup + 1 个 if）。
type RateLimiter struct {
	cfg    RateLimitConfig
	mu     sync.Mutex
	bucket map[string]*tokenBucket
	global *tokenBucket
	now    func() time.Time // 测试注入
}

// RateLimitConfig 是 limiter 的配置。
type RateLimitConfig struct {
	PerTokenPerMinute int
	GlobalPerMinute   int
}

type tokenBucket struct {
	count    int
	windowAt time.Time
}

// NewRateLimiter 用配置构造 limiter；now=nil 时用 time.Now。
func NewRateLimiter(cfg RateLimitConfig, now func() time.Time) *RateLimiter {
	if now == nil {
		now = time.Now
	}
	return &RateLimiter{
		cfg:    cfg,
		bucket: map[string]*tokenBucket{},
		global: &tokenBucket{},
		now:    now,
	}
}

// Allow 返回 (allowed, reason)。reason 仅在 !allowed 时有值，
// 取值 ∈ {"per_token", "global"} 用于错误消息区分。
func (l *RateLimiter) Allow(tokenID string) (bool, string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := l.now()

	// per-token
	b := l.bucket[tokenID]
	if b == nil {
		b = &tokenBucket{}
		l.bucket[tokenID] = b
	}
	if !sameMinute(b.windowAt, now) {
		b.windowAt = now
		b.count = 0
	}
	if b.count >= l.cfg.PerTokenPerMinute {
		return false, "per_token"
	}

	// global
	if !sameMinute(l.global.windowAt, now) {
		l.global.windowAt = now
		l.global.count = 0
	}
	if l.global.count >= l.cfg.GlobalPerMinute {
		return false, "global"
	}

	// 两层桶都通过：增量
	b.count++
	l.global.count++
	return true, ""
}

// sameMinute 判定两时间是否在同一分钟窗口内。
func sameMinute(a, b time.Time) bool {
	return a.Truncate(time.Minute).Equal(b.Truncate(time.Minute))
}

// NewRateLimitMiddleware 必须挂载在 NewAuthMiddleware 之后。
// 从 ctx 取 AuthSession.TokenID 作为 bucket key；找不到 = middleware 顺序错，
// 返 500（避免 fail-open 让无鉴权请求通过 rate limit）。
func NewRateLimitMiddleware(limiter *RateLimiter) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			sess, ok := FromContext(r.Context())
			if !ok {
				httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()),
					errs.New(errs.CodeInternal,
						"rate limit middleware: missing auth session (middleware order bug)"))
				return
			}
			allowed, reason := limiter.Allow(sess.TokenID)
			if !allowed {
				httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()),
					errs.New(errs.CodeMcpRateLimited,
						"rate limit exceeded ("+reason+")"))
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}