package middleware

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

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

var (
	limiters   = make(map[string]*rate.Limiter)
	limiterMux sync.RWMutex
)

// getLimiter retrieves or creates a rate limiter for a specific user.
// Uses a simple token bucket. For strict "10 per day" with distributed persistence,
// this should be refactored to use Redis or DB API usage counters.
func getLimiter(userID string, tier string) *rate.Limiter {
	limiterMux.Lock()
	defer limiterMux.Unlock()

	if limiter, exists := limiters[userID]; exists {
		return limiter
	}

	var limiter *rate.Limiter
	if tier == "Pro" || tier == "Premium" {
		// Unlimited (e.g., 20 requests per second burst)
		limiter = rate.NewLimiter(rate.Limit(20), 100)
	} else {
		// Free: 10 per day -> replenishes 1 token every 2.4 hours, bucket size 10
		limiter = rate.NewLimiter(rate.Every(24*time.Hour/10), 10)
	}

	limiters[userID] = limiter
	return limiter
}

// RateLimit middleware enforces rate limits based on user tier.
// It expects JWTAuth to have already populated UserIDKey in the context.
func RateLimit(db *gorm.DB) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			userIDVal := r.Context().Value(UserIDKey)
			if userIDVal == nil {
				// Allow if not authenticated strictly, or rate limit by IP
				// For now, fallback to generic tight limit for anonymous usage
				ipLimiter := getLimiter(r.RemoteAddr, "Free")
				if !ipLimiter.Allow() {
					log.Printf("[RateLimit] 429 for anonymous IP=%s path=%s", r.RemoteAddr, r.URL.Path)
					http.Error(w, `{"code":429, "message":"Too Many Requests: Rate limit exceeded"}`, http.StatusTooManyRequests)
					return
				}
				next.ServeHTTP(w, r)
				return
			}

			userID := userIDVal.(string)

			// Fast DB query to get user tier (ideally cached in Redis in prod)
			var tier string
			// Look up active subscription for this user
			err := db.Table("subscriptions").
				Select("tier").
				Where("user_id = ? AND status = 'active'", userID).
				Scan(&tier).Error

			if err != nil || tier == "" {
				tier = "Free" // defaults to Free if no active sub
			}

			limiter := getLimiter(userID, tier)
			if !limiter.Allow() {
				log.Printf("[RateLimit] 429 for userID=%s tier=%s path=%s", userID, tier, r.URL.Path)
				http.Error(w, `{"code":429, "message":"Too Many Requests: Daily quota or rate limit exceeded"}`, http.StatusTooManyRequests)
				return
			}

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