// Package middleware holds gin-agnostic-ish HTTP cross-cutting concerns
// (currently: per-IP rate limiting) shared by the v2 gateway router.
package middleware

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

	"github.com/gin-gonic/gin"
)

// bucket 单个 IP 的令牌桶状态。
type bucket struct {
	mu      sync.Mutex
	tokens  float64
	updated time.Time
}

// maxBuckets 是 buckets map 的硬上限——防御性 backstop:哪怕 SetTrustedProxies(nil)
// 已经堵死"伪造 X-Forwarded-For 绕过限流"这条路,socket 对端 IP 本身也可能被
// 大量真实/伪装源地址(如僵尸网络)刷出海量不同 IP,惰性 sweep(见 sweepLocked)
// 只回收超过 10min 不活跃的桶,短时间集中攻击仍可能在 sweep 生效前把 map 撑爆。
// 100k 条 * 单桶几十字节量级,内存占用可控(几 MB),对当前单机部署(阿里云,
// ~1-2GB 内存假设,与 pangolin 单机 Runbook 同量级)足够安全。超限时宁可误限
// (安全侧),也不无界建桶 OOM。
const maxBuckets = 100_000

// IPRateLimiter 简单内存令牌桶,按客户端 IP 分桶限速——单实例部署足够(jiu 反馈波,
// P3 v2 限流纵深)。多实例横向扩展需换外置存储(Redis INCR+EXPIRE 等)统一计数,
// 当前 pay 只跑单进程,内存方案够用且零额外依赖。
type IPRateLimiter struct {
	mu           sync.Mutex
	buckets      map[string]*bucket
	rate         float64 // 每秒补充的令牌数
	burst        float64 // 桶容量(=突发上限,取每分钟额度)
	sweepCounter uint64
}

// NewIPRateLimiter 按每分钟额度构造限流器;perMinute<=0 时按 30 兜底(与 config 里
// "零值=默认开、默认 30/min" 的约定一致,调用方无需重复判空)。
func NewIPRateLimiter(perMinute int) *IPRateLimiter {
	if perMinute <= 0 {
		perMinute = 30
	}
	return &IPRateLimiter{
		buckets: make(map[string]*bucket),
		rate:    float64(perMinute) / 60.0,
		burst:   float64(perMinute),
	}
}

// allow 令牌桶判定:按经过时间匀速补充令牌(封顶 burst),够 1 个则放行并扣 1。
func (l *IPRateLimiter) allow(ip string) bool {
	now := time.Now()
	l.mu.Lock()
	b, ok := l.buckets[ip]
	if !ok {
		if len(l.buckets) >= maxBuckets {
			// map 已到硬上限:先尝试一次 sweep 腾陈旧桶的位置。
			l.sweepLocked(now)
		}
		if len(l.buckets) >= maxBuckets {
			// sweep 后仍超限(短时间内大量活跃桶):拒绝新建桶,按限流命中处理
			// (429)而不是无界增长——安全侧,宁可误限未知新 IP 也不 OOM。
			l.mu.Unlock()
			return false
		}
		b = &bucket{tokens: l.burst, updated: now}
		l.buckets[ip] = b
	}
	l.sweepCounter++
	if l.sweepCounter%1000 == 0 {
		l.sweepLocked(now)
	}
	l.mu.Unlock()

	b.mu.Lock()
	defer b.mu.Unlock()
	elapsed := now.Sub(b.updated).Seconds()
	b.tokens += elapsed * l.rate
	if b.tokens > l.burst {
		b.tokens = l.burst
	}
	b.updated = now
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

// sweepLocked 惰性回收长期不活跃(>10min 无请求)的 IP 桶,防长期运行的单实例
// 在"短连接、海量不同 IP"场景下内存无界增长。调用方须已持 l.mu。
func (l *IPRateLimiter) sweepLocked(now time.Time) {
	for ip, b := range l.buckets {
		b.mu.Lock()
		stale := now.Sub(b.updated) > 10*time.Minute
		b.mu.Unlock()
		if stale {
			delete(l.buckets, ip)
		}
	}
}

// Gin 返回 gin 中间件:超速回 429,不中断其它路由。
func (l *IPRateLimiter) Gin() gin.HandlerFunc {
	return func(c *gin.Context) {
		if !l.allow(c.ClientIP()) {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"code": "rate_limited", "message": "请求过于频繁,请稍后重试",
			})
			return
		}
		c.Next()
	}
}