package middleware

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

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

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

func newIPLimiter(r rate.Limit, burst int) *ipLimiter {
	return &ipLimiter{
		limiters: make(map[string]*rate.Limiter),
		rate:     r,
		burst:    burst,
	}
}

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
	}
	return lim
}

// 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)
		})
	}
}