package middleware

import (
	"RenewCMS/api/pkg/apperr"
	"net/http"
	"strconv"
	"sync"
	"time"
)

type successSlots struct {
	mu     sync.Mutex
	window time.Duration
	now    func() time.Time
	taken  map[string]time.Time
}

func NewUserSuccessLimiter(window time.Duration) func(http.Handler) http.Handler {
	return NewUserSuccessLimiterAt(window, time.Now)
}

func NewUserSuccessLimiterAt(window time.Duration, now func() time.Time) func(http.Handler) http.Handler {
	slots := &successSlots{window: window, now: now, taken: make(map[string]time.Time)}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			key := userRateLimitKey(r)
			claimed, retryAfter, granted := slots.claim(key)
			if !granted {
				w.Header().Set("Retry-After", strconv.Itoa(retryAfter))
				http.Error(w, apperr.TooManyRequestsErrorMessage, http.StatusTooManyRequests)
				return
			}

			recorder := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
			succeeded := false
			defer func() {
				if !succeeded {
					slots.release(key, claimed)
				}
			}()
			next.ServeHTTP(recorder, r)
			succeeded = recorder.status < http.StatusBadRequest
		})
	}
}

func (s *successSlots) claim(key string) (time.Time, int, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := s.now()
	s.sweepExpired(now)
	if claimedAt, held := s.taken[key]; held {
		return time.Time{}, retryAfterSeconds(s.window - now.Sub(claimedAt)), false
	}
	s.taken[key] = now
	return now, 0, true
}

func (s *successSlots) release(key string, claimedAt time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if current, held := s.taken[key]; held && current.Equal(claimedAt) {
		delete(s.taken, key)
	}
}

func (s *successSlots) sweepExpired(now time.Time) {
	for key, claimedAt := range s.taken {
		if now.Sub(claimedAt) >= s.window {
			delete(s.taken, key)
		}
	}
}

type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (s *statusRecorder) WriteHeader(status int) {
	s.status = status
	s.ResponseWriter.WriteHeader(status)
}

func (s *statusRecorder) Unwrap() http.ResponseWriter {
	return s.ResponseWriter
}