package middleware

import (
	"sync"
	"time"

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

type RateLimitConfig struct {
	Max        int
	WindowSecs int
	KeyFunc    func(*fiber.Ctx) string
}

type rateLimiter struct {
	requests map[string][]time.Time
	mu       sync.RWMutex
	max      int
	window   time.Duration
}

func newRateLimiter(max int, windowSecs int) *rateLimiter {
	rl := &rateLimiter{
		requests: make(map[string][]time.Time),
		max:      max,
		window:   time.Duration(windowSecs) * time.Second,
	}

	go rl.cleanup()
	return rl
}

func (rl *rateLimiter) cleanup() {
	ticker := time.NewTicker(time.Minute)
	for range ticker.C {
		rl.mu.Lock()
		now := time.Now()
		for key, times := range rl.requests {
			var valid []time.Time
			for _, t := range times {
				if now.Sub(t) < rl.window {
					valid = append(valid, t)
				}
			}
			if len(valid) == 0 {
				delete(rl.requests, key)
			} else {
				rl.requests[key] = valid
			}
		}
		rl.mu.Unlock()
	}
}

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

	now := time.Now()
	windowStart := now.Add(-rl.window)

	times := rl.requests[key]
	var valid []time.Time
	for _, t := range times {
		if t.After(windowStart) {
			valid = append(valid, t)
		}
	}

	if len(valid) >= rl.max {
		rl.requests[key] = valid
		return false
	}

	rl.requests[key] = append(valid, now)
	return true
}

func RateLimit(config RateLimitConfig) fiber.Handler {
	if config.Max == 0 {
		config.Max = 100
	}
	if config.WindowSecs == 0 {
		config.WindowSecs = 60
	}
	if config.KeyFunc == nil {
		config.KeyFunc = func(c *fiber.Ctx) string {
			return c.IP()
		}
	}

	limiter := newRateLimiter(config.Max, config.WindowSecs)

	return func(c *fiber.Ctx) error {
		key := config.KeyFunc(c)

		if !limiter.allow(key) {
			return c.Status(429).JSON(fiber.Map{
				"error": "Too many requests",
			})
		}

		return c.Next()
	}
}