package middleware

import (
	"sync"
	"time"

	"github.com/gofiber/fiber/v2"
)

type bucket struct {
	tokens    float64
	lastCheck time.Time
}

type rateLimiter struct {
	mu       sync.Mutex
	buckets  map[string]*bucket
	rate     float64 // tokens per second
	capacity float64
}

func NewRateLimiter(requestsPerMinute int) fiber.Handler {
	rl := &rateLimiter{
		buckets:  make(map[string]*bucket),
		rate:     float64(requestsPerMinute) / 60.0,
		capacity: float64(requestsPerMinute),
	}

	return func(c *fiber.Ctx) error {
		ip := c.IP()
		rl.mu.Lock()
		b, ok := rl.buckets[ip]
		if !ok {
			b = &bucket{tokens: rl.capacity, lastCheck: time.Now()}
			rl.buckets[ip] = b
		}

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

		if b.tokens < 1 {
			rl.mu.Unlock()
			return c.Status(429).JSON(fiber.Map{"error": "rate limit exceeded"})
		}
		b.tokens--
		rl.mu.Unlock()

		return c.Next()
	}
}