package middleware

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

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

type IPRateLimiter struct {
	mu       sync.RWMutex
	limiters map[string]*rate.Limiter
	r        rate.Limit
	b        int
}

func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
	return &IPRateLimiter{
		limiters: make(map[string]*rate.Limiter),
		r:        r,
		b:        b,
	}
}

func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
	i.mu.Lock()
	defer i.mu.Unlock()
	if lim, ok := i.limiters[ip]; ok {
		return lim
	}
	lim := rate.NewLimiter(i.r, i.b)
	i.limiters[ip] = lim
	return lim
}

func (i *IPRateLimiter) Cleanup() {
	ticker := time.NewTicker(time.Minute)
	go func() {
		for range ticker.C {
			i.mu.Lock()
			for k := range i.limiters {
				delete(i.limiters, k)
			}
			i.mu.Unlock()
		}
	}()
}

func RateLimit(l *IPRateLimiter) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ip := r.RemoteAddr
			if xf := r.Header.Get("X-Forwarded-For"); xf != "" {
				// simples, suficiente pra começar
				ip = xf
			}
			if !l.GetLimiter(ip).Allow() {
				w.WriteHeader(http.StatusTooManyRequests)
				w.Write([]byte(`{"error":"muitas tentativas, tente novamente mais tarde"}`))
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}