package middleware

import (
	"sync"
	"time"
	"yintai-api/base/config"

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

func LoginLimit() gin.HandlerFunc {
	const maxPerSecond = 2
	var (
		mu   sync.Mutex
		hits = map[string][]time.Time{}
	)
	return func(c *gin.Context) {
		ip := c.ClientIP()
		now := time.Now()
		cutoff := now.Add(-time.Second)

		mu.Lock()
		recent := make([]time.Time, 0, maxPerSecond)
		for _, t := range hits[ip] {
			if t.After(cutoff) {
				recent = append(recent, t)
			}
		}
		if len(recent) >= maxPerSecond {
			hits[ip] = recent
			mu.Unlock()
			panic(config.ErrLimit.New())
		}
		hits[ip] = append(recent, now)
		if len(hits) > 1024 {
			for k, v := range hits {
				alive := false
				for _, t := range v {
					if t.After(cutoff) {
						alive = true
						break
					}
				}
				if !alive {
					delete(hits, k)
				}
			}
		}
		mu.Unlock()
		c.Next()
	}
}