package httpserver

import (
	"math"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

type rateLimitStore struct {
	mu          sync.Mutex
	ratePerSec  float64
	burst       float64
	lastCleanup time.Time
	entries     map[string]*rateLimitEntry
}

type rateLimitEntry struct {
	tokens   float64
	lastSeen time.Time
}

func newRateLimitStore(perMinute, burst int) *rateLimitStore {
	return &rateLimitStore{
		ratePerSec: float64(perMinute) / 60,
		burst:      float64(burst),
		entries:    make(map[string]*rateLimitEntry),
	}
}

func (s *rateLimitStore) Allow(key string, now time.Time) (bool, time.Duration) {
	key = strings.TrimSpace(key)
	if key == "" {
		key = "global"
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	if entry, ok := s.entries[key]; ok {
		elapsed := now.Sub(entry.lastSeen).Seconds()
		entry.tokens = math.Min(s.burst, entry.tokens+(elapsed*s.ratePerSec))
		entry.lastSeen = now
		if entry.tokens >= 1 {
			entry.tokens--
			s.cleanup(now)
			return true, 0
		}

		s.cleanup(now)
		return false, retryAfter(entry.tokens, s.ratePerSec)
	}

	s.entries[key] = &rateLimitEntry{
		tokens:   s.burst - 1,
		lastSeen: now,
	}
	s.cleanup(now)
	return true, 0
}

func (s *rateLimitStore) cleanup(now time.Time) {
	if !s.lastCleanup.IsZero() && now.Sub(s.lastCleanup) < 5*time.Minute {
		return
	}

	for key, entry := range s.entries {
		if now.Sub(entry.lastSeen) > 15*time.Minute {
			delete(s.entries, key)
		}
	}

	s.lastCleanup = now
}

func retryAfter(tokens, ratePerSec float64) time.Duration {
	if ratePerSec <= 0 {
		return time.Minute
	}

	if tokens < 0 {
		tokens = 0
	}

	seconds := (1 - tokens) / ratePerSec
	if seconds < 1 {
		seconds = 1
	}

	return time.Duration(math.Ceil(seconds * float64(time.Second)))
}

func retryAfterHeader(delay time.Duration) string {
	seconds := int(math.Ceil(delay.Seconds()))
	if seconds < 1 {
		seconds = 1
	}

	return strconv.Itoa(seconds)
}

func (h *handler) loginRateLimit(next http.Handler) http.Handler {
	return h.rateLimit(h.loginRateLimiter, func(r *http.Request) string {
		return clientIP(r)
	}, func(w http.ResponseWriter, r *http.Request, delay time.Duration) {
		w.Header().Set("Retry-After", retryAfterHeader(delay))
		http.Error(w, "too many login attempts", http.StatusTooManyRequests)
	})(next)
}

func (h *handler) clientAPIRateLimit(next http.Handler) http.Handler {
	return h.rateLimit(h.clientRateLimiter, func(r *http.Request) string {
		return clientIP(r)
	}, func(w http.ResponseWriter, r *http.Request, delay time.Duration) {
		w.Header().Set("Retry-After", retryAfterHeader(delay))
		writeJSON(w, http.StatusTooManyRequests, map[string]any{"error": "rate limit exceeded"})
	})(next)
}

func (h *handler) rateLimit(store *rateLimitStore, keyFn func(*http.Request) string, reject func(http.ResponseWriter, *http.Request, time.Duration)) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			key := clientIP(r)
			if keyFn != nil {
				key = keyFn(r)
			}

			allowed, delay := store.Allow(key, time.Now().UTC())
			if !allowed {
				reject(w, r, delay)
				return
			}

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