package middleware

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

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

type RateLimiter struct {
	mu       sync.Mutex
	clients  map[string]*clientBuckets
	rate     int
	burst    int
}

type clientBuckets struct {
	tokens   int
	lastFill time.Time
}

func NewRateLimiter(rate, burst int) *RateLimiter {
	return &RateLimiter{
		clients: make(map[string]*clientBuckets),
		rate:    rate,
		burst:   burst,
	}
}

func (rl *RateLimiter) Handler() gin.HandlerFunc {
	return func(c *gin.Context) {
		if !rl.allow(c.ClientIP()) {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "rate limit exceeded"})
			return
		}
		c.Next()
	}
}

func (rl *RateLimiter) allow(key string) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	b, ok := rl.clients[key]
	if !ok {
		b = &clientBuckets{tokens: rl.burst, lastFill: time.Now()}
		rl.clients[key] = b
	}

	now := time.Now()
	elapsed := now.Sub(b.lastFill)
	b.lastFill = now
	b.tokens += int(elapsed.Seconds()) * rl.rate
	if b.tokens > rl.burst {
		b.tokens = rl.burst
	}

	if b.tokens > 0 {
		b.tokens--
		return true
	}
	return false
}