gochat/backend/internal/middleware/rate_limit.go

package middleware

import (
	"context"
	"fmt"
	"net/http"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"

	"github.com/gochat/gochat/pkg/logger"
	"github.com/gochat/gochat/pkg/response"
)

const (
	// rateLimitKeyPrefix is the Redis key prefix for rate limit counters.
	rateLimitKeyPrefix = "gochat:rate_limit:"

	// Hardcoded rate limit defaults (no longer configurable).
	rateLimitRequestsPerMinute = 100
	rateLimitWindowSeconds     = 60
)

// slidingWindowLimiter implements Redis-based sliding window counter rate limiting.
// Falls back to in-memory limiting when Redis is unavailable.
// Reference: Chatwoot's Rack::Attack throttle configuration.
type slidingWindowLimiter struct {
	redis          *redis.Client
	fallback       *inMemoryLimiter
	redisAvailable atomic.Bool
}

// inMemoryLimiter provides in-memory rate limiting as a fallback when Redis is unavailable.
type inMemoryLimiter struct {
	mu       sync.Mutex
	visitors map[string]*visitorEntry
	limit    int
	window   time.Duration
}

type visitorEntry struct {
	count    int
	lastSeen time.Time
}

// newInMemoryLimiter creates a new in-memory rate limiter with the given limit and window.
func newInMemoryLimiter(limit int, window time.Duration) *inMemoryLimiter {
	im := &inMemoryLimiter{
		visitors: make(map[string]*visitorEntry),
		limit:    limit,
		window:   window,
	}

	// Cleanup old entries periodically
	go func() {
		for {
			time.Sleep(window)
			im.mu.Lock()
			for key, v := range im.visitors {
				if time.Since(v.lastSeen) > window {
					delete(im.visitors, key)
				}
			}
			im.mu.Unlock()
		}
	}()

	return im
}

// checkInMemory checks rate limit using in-memory store. Returns (allowed, currentCount).
func (im *inMemoryLimiter) checkInMemory(key string) (bool, int) {
	im.mu.Lock()
	defer im.mu.Unlock()

	v, exists := im.visitors[key]
	if !exists {
		im.visitors[key] = &visitorEntry{count: 1, lastSeen: time.Now()}
		return true, 1
	}

	if time.Since(v.lastSeen) > im.window {
		v.count = 1
		v.lastSeen = time.Now()
		return true, 1
	}

	v.count++
	v.lastSeen = time.Now()

	if v.count > im.limit {
		return false, v.count
	}
	return true, v.count
}

// remainingInMemory returns remaining request count from in-memory store.
func (im *inMemoryLimiter) remainingInMemory(key string) int {
	im.mu.Lock()
	defer im.mu.Unlock()

	v, exists := im.visitors[key]
	if !exists {
		return im.limit
	}
	if time.Since(v.lastSeen) > im.window {
		return im.limit
	}
	remaining := im.limit - v.count
	if remaining < 0 {
		return 0
	}
	return remaining
}

// newSlidingWindowLimiter creates a new rate limiter with Redis sliding window counter
// and in-memory fallback.
func newSlidingWindowLimiter(rdb *redis.Client) *slidingWindowLimiter {
	limit := rateLimitRequestsPerMinute
	windowSecs := rateLimitWindowSeconds

	sw := &slidingWindowLimiter{
		redis:    rdb,
		fallback: newInMemoryLimiter(limit, time.Duration(windowSecs)*time.Second),
	}

	// Initial Redis availability check
	if rdb != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		if err := rdb.Ping(ctx).Err(); err != nil {
			sw.redisAvailable.Store(false)
			logger.L().Warnf("Redis unavailable for rate limiting, falling back to in-memory: %v", err)
		} else {
			sw.redisAvailable.Store(true)
		}
	} else {
		sw.redisAvailable.Store(false)
		logger.L().Warn("No Redis client provided for rate limiting, using in-memory fallback")
	}

	return sw
}

// checkRedis performs a Redis sliding window counter check.
// Returns (allowed, currentCount, remaining, error).
//
// The sliding window counter algorithm:
// 1. Use two fixed windows: current window and previous window
// 2. Calculate weighted count from previous window based on elapsed time ratio
// 3. Total count = previous window weighted count + current window count
// 4. If total count > limit, reject; otherwise allow and increment current window
func (sw *slidingWindowLimiter) checkRedis(ctx context.Context, key string) (bool, int, int, error) {
	now := time.Now()
	windowSecs := rateLimitWindowSeconds
	limit := rateLimitRequestsPerMinute

	currentWindow := now.Unix() / int64(windowSecs)
	previousWindow := currentWindow - 1

	currentKey := fmt.Sprintf("%s%s:%d", rateLimitKeyPrefix, key, currentWindow)
	previousKey := fmt.Sprintf("%s%s:%d", rateLimitKeyPrefix, key, previousWindow)

	// Use a Lua script for atomicity: calculate sliding window count + increment
	luaScript := redis.NewScript(`
local current_key = KEYS[1]
local previous_key = KEYS[2]
local limit = tonumber(ARGV[1])
local window_secs = tonumber(ARGV[2])
local elapsed_ratio = tonumber(ARGV[3])

-- Get previous window count (weighted by remaining time)
local previous_count = tonumber(redis.call("GET", previous_key) or "0")
local weighted_previous = math.floor(previous_count * elapsed_ratio)

-- Get current window count
local current_count = tonumber(redis.call("GET", current_key) or "0")

-- Calculate total count in sliding window
local total = weighted_previous + current_count

-- Check if over limit
if total >= limit then
    return {total, 0}
end

-- Increment current window
local new_count = redis.call("INCR", current_key)
if new_count == 1 then
    redis.call("EXPIRE", current_key, window_secs * 2)
end

-- Recalculate total with the new increment
local new_total = weighted_previous + new_count
local remaining = limit - new_total
if remaining < 0 then
    remaining = 0
end

return {new_total, remaining}
`)

	// Calculate elapsed ratio: fraction of the current window that has elapsed
	elapsedInWindow := now.Unix() % int64(windowSecs)
	elapsedRatio := float64(windowSecs-int(elapsedInWindow)) / float64(windowSecs)

	result, err := luaScript.Run(ctx, sw.redis, []string{currentKey, previousKey},
		limit, windowSecs, elapsedRatio).Int64Slice()
	if err != nil {
		return false, 0, 0, fmt.Errorf("redis sliding window script: %w", err)
	}

	totalCount := int(result[0])
	remaining := int(result[1])

	allowed := totalCount <= limit
	return allowed, totalCount, remaining, nil
}

// check performs rate limiting, using Redis when available and falling back to in-memory.
// Returns (allowed, currentCount, remaining, usedRedis).
func (sw *slidingWindowLimiter) check(ctx context.Context, key string) (bool, int, int, bool) {
	if sw.redisAvailable.Load() && sw.redis != nil {
		allowed, count, remaining, err := sw.checkRedis(ctx, key)
		if err != nil {
			// Redis error — mark unavailable and fall back to in-memory
			sw.redisAvailable.Store(false)
			logger.L().Warnf("Redis rate limit check failed, falling back to in-memory: %v", err)

			// Schedule Redis availability recheck
			go sw.recheckRedis()

			allowed, count = sw.fallback.checkInMemory(key)
			remaining = sw.fallback.remainingInMemory(key)
			return allowed, count, remaining, false
		}
		return allowed, count, remaining, true
	}

	// In-memory fallback
	allowed, count := sw.fallback.checkInMemory(key)
	remaining := sw.fallback.remainingInMemory(key)
	return allowed, count, remaining, false
}

// recheckRedis attempts to re-establish Redis connectivity for rate limiting.
func (sw *slidingWindowLimiter) recheckRedis() {
	time.Sleep(5 * time.Second) // wait before rechecking
	if sw.redis == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	if err := sw.redis.Ping(ctx).Err(); err == nil {
		sw.redisAvailable.Store(true)
		logger.L().Info("Redis rate limiting restored")
	}
}

// RateLimit creates a Redis-based sliding window counter rate limiting middleware.
// Falls back to in-memory rate limiting when Redis is unavailable.
// Corresponds to Chatwoot's Rack::Attack throttle configuration.
func RateLimit(rdb *redis.Client) gin.HandlerFunc {
	sw := newSlidingWindowLimiter(rdb)
	limit := rateLimitRequestsPerMinute

	return func(c *gin.Context) {
		if isRateLimitExemptPath(c.Request.URL.Path) {
			c.Next()
			return
		}

		ip := c.ClientIP()
		key := "global:" + ip

		allowed, _, remaining, usedRedis := sw.check(c.Request.Context(), key)
		// Set rate limit headers
		c.Header("X-RateLimit-Limit", strconv.Itoa(limit))
		c.Header("X-RateLimit-Remaining", strconv.Itoa(remaining))
		if usedRedis {
			c.Header("X-RateLimit-Backend", "redis")
		} else {
			c.Header("X-RateLimit-Backend", "memory")
		}

		if !allowed {
			c.Header("Retry-After", strconv.Itoa(rateLimitWindowSeconds))
			c.AbortWithStatusJSON(http.StatusTooManyRequests, response.APIResponse{
				Success: false,
				Error: &response.ErrorBody{
					Code:    response.ErrRateLimit,
					Message: "Rate limit exceeded",
				},
			})
			return
		}

		c.Next()
	}
}

func isRateLimitExemptPath(path string) bool {
	switch path {
	case "/health", "/ready", "/live", "/metrics":
		return true
	default:
		return false
	}
}

// PerRouteLimit creates a per-route rate limiting middleware.
// This allows different rate limits for different API endpoints, matching
// Chatwoot's Rack::Attack per-route throttle configuration.
//
// Usage:
//
//	router.GET("/conversations", PerRouteLimit("conversations_list", 60), listConversations)
//	router.POST("/messages", PerRouteLimit("messages_create", 30), createMessage)
//	router.GET("/reports", PerRouteLimit("reports_read", 10), viewReports)
//
// The route identifier is used as a key namespace so that limits on one route
// don't affect limits on another route for the same IP.
func PerRouteLimit(route string, requestsPerMin int) gin.HandlerFunc {
	type visitor struct {
		count    int
		lastSeen time.Time
	}
	var (
		mu       sync.Mutex
		visitors = make(map[string]*visitor) // key: route:ip
		limit    = requestsPerMin
		window   = time.Minute
	)

	// Cleanup old entries periodically
	go func() {
		for {
			time.Sleep(window)
			mu.Lock()
			for key, v := range visitors {
				if time.Since(v.lastSeen) > window {
					delete(visitors, key)
				}
			}
			mu.Unlock()
		}
	}()

	return func(c *gin.Context) {
		ip := c.ClientIP()
		key := route + ":" + ip

		mu.Lock()
		v, exists := visitors[key]
		if !exists {
			visitors[key] = &visitor{count: 1, lastSeen: time.Now()}
			mu.Unlock()
			c.Next()
			return
		}

		if time.Since(v.lastSeen) > window {
			v.count = 1
			v.lastSeen = time.Now()
			mu.Unlock()
			c.Next()
			return
		}

		v.count++
		v.lastSeen = time.Now()
		mu.Unlock()

		if v.count > limit {
			c.Header("X-RateLimit-Limit", strconv.Itoa(limit))
			c.Header("X-RateLimit-Remaining", "0")
			c.Header("Retry-After", "60")
			c.AbortWithStatusJSON(http.StatusTooManyRequests, response.APIResponse{
				Success: false,
				Error: &response.ErrorBody{
					Code:    response.ErrRateLimit,
					Message: "Rate limit exceeded for route '" + route + "'. Maximum " + strconv.Itoa(limit) + " requests per minute.",
				},
			})
			return
		}

		// Set rate limit headers for successful requests
		remaining := limit - v.count
		if remaining < 0 {
			remaining = 0
		}
		c.Header("X-RateLimit-Limit", strconv.Itoa(limit))
		c.Header("X-RateLimit-Remaining", strconv.Itoa(remaining))

		c.Next()
	}
}

// PerUserLimit creates a per-user rate limiting middleware that uses the authenticated
// user ID instead of IP address. This provides more accurate limits for authenticated
// endpoints where multiple users may share the same IP (e.g., office networks).
//
// Usage:
//
//	router.POST("/api/v1/conversations", AuthRequired(jwtSvc), PerUserLimit("conversations_create", 30), createConversation)
func PerUserLimit(route string, requestsPerMin int) gin.HandlerFunc {
	type visitor struct {
		count    int
		lastSeen time.Time
	}
	var (
		mu       sync.Mutex
		visitors = make(map[string]*visitor) // key: route:user_id
		limit    = requestsPerMin
		window   = time.Minute
	)

	go func() {
		for {
			time.Sleep(window)
			mu.Lock()
			for key, v := range visitors {
				if time.Since(v.lastSeen) > window {
					delete(visitors, key)
				}
			}
			mu.Unlock()
		}
	}()

	return func(c *gin.Context) {
		// Use user_id from context (set by AuthRequired middleware)
		userID, exists := c.Get("user_id")
		if !exists {
			// Fall back to IP-based limiting if not authenticated
			ip := c.ClientIP()
			key := route + ":ip:" + ip

			mu.Lock()
			v, vExists := visitors[key]
			if !vExists {
				visitors[key] = &visitor{count: 1, lastSeen: time.Now()}
				mu.Unlock()
				c.Next()
				return
			}

			if time.Since(v.lastSeen) > window {
				v.count = 1
				v.lastSeen = time.Now()
				mu.Unlock()
				c.Next()
				return
			}

			v.count++
			v.lastSeen = time.Now()
			mu.Unlock()

			if v.count > limit {
				c.AbortWithStatusJSON(http.StatusTooManyRequests, response.APIResponse{
					Success: false,
					Error: &response.ErrorBody{
						Code:    response.ErrRateLimit,
						Message: "Rate limit exceeded for route '" + route + "'",
					},
				})
				return
			}

			c.Next()
			return
		}

		// User-based rate limiting
		key := route + ":uid:" + strconv.FormatUint(uint64(userID.(uint)), 10)

		mu.Lock()
		v, exists := visitors[key]
		if !exists {
			visitors[key] = &visitor{count: 1, lastSeen: time.Now()}
			mu.Unlock()
			c.Next()
			return
		}

		if time.Since(v.lastSeen) > window {
			v.count = 1
			v.lastSeen = time.Now()
			mu.Unlock()
			c.Next()
			return
		}

		v.count++
		v.lastSeen = time.Now()
		mu.Unlock()

		if v.count > limit {
			c.Header("X-RateLimit-Limit", strconv.Itoa(limit))
			c.Header("X-RateLimit-Remaining", "0")
			c.AbortWithStatusJSON(http.StatusTooManyRequests, response.APIResponse{
				Success: false,
				Error: &response.ErrorBody{
					Code:    response.ErrRateLimit,
					Message: "Rate limit exceeded for route '" + route + "'. Maximum " + strconv.Itoa(limit) + " requests per minute per user.",
				},
			})
			return
		}

		remaining := limit - v.count
		if remaining < 0 {
			remaining = 0
		}
		c.Header("X-RateLimit-Limit", strconv.Itoa(limit))
		c.Header("X-RateLimit-Remaining", strconv.Itoa(remaining))

		c.Next()
	}
}