package middleware

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

	"golang.org/x/time/rate"
)

// ipLimiter tracks rate limiters per IP address.
type ipLimiter struct {
	mu       sync.Mutex
	limiters map[string]*rate.Limiter
	lastUsed map[string]time.Time
	rate     rate.Limit
	burst    int
}

func newIPLimiter(r rate.Limit, burst int) *ipLimiter {
	l := &ipLimiter{
		limiters: make(map[string]*rate.Limiter),
		lastUsed: make(map[string]time.Time),
		rate:     r,
		burst:    burst,
	}
	// ponytail: evict inactive limiters every 5 minutes to prevent unbounded map growth
	go l.cleanup()
	return l
}

func (l *ipLimiter) getLimiter(ip string) *rate.Limiter {
	l.mu.Lock()
	defer l.mu.Unlock()

	lim, exists := l.limiters[ip]
	if !exists {
		lim = rate.NewLimiter(l.rate, l.burst)
		l.limiters[ip] = lim
	}
	l.lastUsed[ip] = time.Now()
	return lim
}

// cleanup runs in a background goroutine and removes limiters that haven't
// been accessed in 10 minutes.
func (l *ipLimiter) cleanup() {
	for {
		time.Sleep(5 * time.Minute)
		l.mu.Lock()
		now := time.Now()
		for ip, last := range l.lastUsed {
			if now.Sub(last) > 10*time.Minute {
				delete(l.limiters, ip)
				delete(l.lastUsed, ip)
			}
		}
		l.mu.Unlock()
	}
}

// RateLimit returns middleware that limits requests per IP.
// requestsPerSecond is the sustained rate, burst is the max burst size.
func RateLimit(requestsPerSecond float64, burst int) func(http.Handler) http.Handler {
	limiter := newIPLimiter(rate.Limit(requestsPerSecond), burst)

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ip := r.RemoteAddr
			if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
				// Take the first IP in the chain (the original client)
				if idx := strings.IndexByte(xff, ','); idx > 0 {
					ip = strings.TrimSpace(xff[:idx])
				} else {
					ip = strings.TrimSpace(xff)
				}
			} else if xri := r.Header.Get("X-Real-IP"); xri != "" {
				ip = strings.TrimSpace(xri)
			}

			if !limiter.getLimiter(ip).Allow() {
				http.Error(w, `{"error":"rate limited"}`, http.StatusTooManyRequests)
				return
			}

			next.ServeHTTP(w, r)
		})
	}
}