package middleware

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

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

// RateLimiter provides per-IP rate limiting middleware.
type RateLimiter struct {
	limiters map[string]*rate.Limiter
	mu       sync.RWMutex
	rate     rate.Limit // Requests per second
	burst    int        // Max burst size
}

// NewRateLimiter creates a new rate limiter with the specified rate (requests/second) and burst size.
// Includes a background cleaner to remove stale entries and prevent memory leaks.
func NewRateLimiter(r float64, burst int) *RateLimiter {
	rl := &RateLimiter{
		limiters: make(map[string]*rate.Limiter),
		rate:     rate.Limit(r),
		burst:    burst,
	}

	// Start background cleanup
	go rl.cleanupLoop()
	return rl
}

// cleanupLoop periodically removes limiters that haven't been used recently
// Clears the map every 10 minutes if it grows too large (>10,000 entries)
func (rl *RateLimiter) cleanupLoop() {
	ticker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()

	for range ticker.C {
		rl.mu.Lock()
		if len(rl.limiters) > 10000 {
			// Reset the map - limiters are cheap to recreate
			rl.limiters = make(map[string]*rate.Limiter)
		}
		rl.mu.Unlock()
	}
}

// getLimiter retrieves or creates a rate limiter for the given key (typically IP address).
func (rl *RateLimiter) getLimiter(key string) *rate.Limiter {
	rl.mu.RLock()
	limiter, exists := rl.limiters[key]
	rl.mu.RUnlock()

	if exists {
		return limiter
	}

	rl.mu.Lock()
	defer rl.mu.Unlock()

	// Double-check after acquiring write lock
	if limiter, exists = rl.limiters[key]; exists {
		return limiter
	}

	limiter = rate.NewLimiter(rl.rate, rl.burst)
	rl.limiters[key] = limiter
	return limiter
}

// Middleware returns a Gin middleware that applies rate limiting.
// Bypasses the /health endpoint.
func (rl *RateLimiter) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		// Bypass health checks
		if c.Request.URL.Path == "/health" {
			c.Next()
			return
		}

		// Use client IP as the rate limit key
		clientIP := c.ClientIP()
		limiter := rl.getLimiter(clientIP)

		if !limiter.Allow() {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"error":       "Too Many Requests",
				"message":     "Rate limit exceeded. Please slow down.",
				"retry_after": "1s",
			})
			return
		}

		c.Next()
	}
}

// DefaultRateLimiter creates a limiter with configurable or default values.
// Environment variables: RATE_LIMIT_RPS (default: 100), RATE_LIMIT_BURST (default: 200)
func DefaultRateLimiter() *RateLimiter {
	rps := getEnvFloat("RATE_LIMIT_RPS", 100)
	burst := getEnvInt("RATE_LIMIT_BURST", 200)
	return NewRateLimiter(rps, burst)
}

// getEnvFloat reads an environment variable as float64 with fallback
func getEnvFloat(key string, fallback float64) float64 {
	if val := os.Getenv(key); val != "" {
		if f, err := strconv.ParseFloat(val, 64); err == nil {
			return f
		}
	}
	return fallback
}

// getEnvInt reads an environment variable as int with fallback
func getEnvInt(key string, fallback int) int {
	if val := os.Getenv(key); val != "" {
		if i, err := strconv.Atoi(val); err == nil {
			return i
		}
	}
	return fallback
}