package middleware

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

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

type RateLimiter struct {
	visitors map[string]*Visitor
	mu       sync.Mutex
}

type Visitor struct {
	lastSeen time.Time
	tokens   int
}

func NewRateLimiter() *RateLimiter {
	rl := &RateLimiter{
		visitors: make(map[string]*Visitor),
	}
	// 定期清理过期访问者
	go rl.cleanupVisitors()
	return rl
}

func (rl *RateLimiter) cleanupVisitors() {
	for {
		time.Sleep(time.Minute)
		rl.mu.Lock()
		for ip, v := range rl.visitors {
			if time.Since(v.lastSeen) > time.Minute {
				delete(rl.visitors, ip)
			}
		}
		rl.mu.Unlock()
	}
}

func (rl *RateLimiter) LimitMiddleware(maxRequests int, per time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		ip := c.ClientIP()
		rl.mu.Lock()
		defer rl.mu.Unlock()

		v, exists := rl.visitors[ip]
		if !exists {
			rl.visitors[ip] = &Visitor{
				lastSeen: time.Now(),
				tokens:   maxRequests - 1,
			}
			c.Next()
			return
		}

		// 重置 token
		if time.Since(v.lastSeen) > per {
			v.tokens = maxRequests - 1
			v.lastSeen = time.Now()
			c.Next()
			return
		}

		if v.tokens > 0 {
			v.tokens--
			v.lastSeen = time.Now()
			c.Next()
			return
		}

		// 超过限流
		c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
			"error": "请求过于频繁，请稍后再试",
		})
	}
}