// Package ratelimit provides an HTTP middleware that rate-limits requests
// on a per-client basis using the token bucket algorithm.
//
// Ecample:
//
//	limiter := ratelimit.MiddlewareWithConfig(ratelimit.Config{
//	    RPS:   10,
//	    Burst: 20,
//	    TTL:   time.Minute,
//	})
//	http.Handle("/api", limiter(yourHandler))
//
package ratelimit

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

	"olexsmir.xyz/x/ratelimit/internal/time/rate"
)

type (
	rateLimiter struct {
		mu sync.RWMutex

		visitors map[visitorIP]*visitor

		// limit is the maximum number of requests per second
		limit rate.Limit

		// ttl is the time after which a visitor is forgotten
		ttl time.Duration

		// burst is the maximum number of requests allowed in a short burst
		burst int
	}

	visitorIP string
	visitor   struct {
		limiter  *rate.Limiter
		lastSeen time.Time
	}
)

func newLimiter(rps, burst int, ttl time.Duration) *rateLimiter {
	return &rateLimiter{
		visitors: make(map[visitorIP]*visitor),
		limit:    rate.Limit(rps),
		burst:    burst,
		ttl:      ttl,
	}
}

// getVisitor retrieves the rate limiter for the given IP, or creates one.
func (r *rateLimiter) getVisitor(ip visitorIP) *rate.Limiter {
	r.mu.RLock()
	v, exists := r.visitors[ip]
	r.mu.RUnlock()

	if !exists {
		limit := rate.NewLimiter(r.limit, r.burst)

		r.mu.Lock()
		r.visitors[ip] = &visitor{
			limiter:  limit,
			lastSeen: time.Now(),
		}
		r.mu.Unlock()

		return limit
	}

	r.mu.Lock()
	v.lastSeen = time.Now()
	r.mu.Unlock()

	return v.limiter
}

// cleanupVisitors removes visitors that haven't been seen for longer than TTL.
func (r *rateLimiter) cleanupVisitors() {
	r.mu.Lock()
	defer r.mu.Unlock()

	for ip, v := range r.visitors {
		if time.Since(v.lastSeen) > r.ttl {
			delete(r.visitors, ip)
		}
	}
}

// cleanupVisitorsLoop runs cleanupVisitors every minute.
func (r *rateLimiter) cleanupVisitorsLoop() {
	for {
		time.Sleep(time.Minute)
		r.cleanupVisitors()
	}
}

// Config configures the rate limiter middleware.
type Config struct {
	// RPS is the maximum number of requests per second per visitor.
	RPS int

	// TTL is the duration after which an idle visitor's limiter is cleaned up.
	TTL time.Duration

	// Burst is the maximum number of requests that can be made in a short burst.
	Burst int
}

// Middleware returns a rate-limiting HTTP middleware.
//
// It limits requests per client IP using the token bucket algorithm.
// When the limit is exceeded it responds with 429 Too Many Requests.
func Middleware(c Config) func(http.Handler) http.Handler {
	lmt := newLimiter(c.RPS, c.Burst, c.TTL)
	go lmt.cleanupVisitorsLoop()

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ip := getIP(r)
			limiter := lmt.getVisitor(ip)
			if limiter == nil {
				http.Error(w, http.StatusText(http.StatusInternalServerError), http.StatusInternalServerError)
				return
			}

			if !limiter.Allow() {
				http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
				return
			}

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

func getIP(r *http.Request) visitorIP {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		if i := strings.IndexByte(xff, ','); i != -1 {
			xff = xff[:i]
		}
		if ip := strings.TrimSpace(xff); ip != "" {
			return visitorIP(ip)
		}
	}

	if xri := r.Header.Get("X-Real-IP"); xri != "" {
		return visitorIP(xri)
	}

	if host, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		return visitorIP(host)
	}

	return visitorIP(r.RemoteAddr)
}