package middlewares

import (
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/sethvargo/go-limiter"
)

type RateLimiter struct {
	store limiter.Store
	next  http.Handler
}

func (rl *RateLimiter) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// What to not rate limit
	// Includes various localhost and loopback interfaces
	// TODO: Only allow requests with a valid unlimit token in the "rate-limit-bypass" field bypass rate limit
	if r.FormValue("rate-limit-bypass") != "" ||
		strings.Contains(r.RemoteAddr, "127.0.0.1") ||
		strings.Contains(r.RemoteAddr, "localhost") ||
		strings.Contains(r.RemoteAddr, "::1") {
		rl.next.ServeHTTP(w, r)
		return
	}
	tokens, remaining, resetTime, ok, err := rl.store.Take(r.Context(), r.RemoteAddr)
	if err != nil {
		http.Error(w, "rate limiter problem", http.StatusInternalServerError)
		return
	}
	w.Header().Add("X-RateLimit-Limit", fmt.Sprint(tokens))
	w.Header().Add("X-RateLimit-Remaining", fmt.Sprint(remaining))
	w.Header().Add("X-RateLimit-Reset", fmt.Sprint(resetTime))
	if ok {
		rl.next.ServeHTTP(w, r)
	} else {
		t := time.Unix(0, int64(resetTime)).UTC().Format(time.RFC1123)
		w.Header().Add("Retry-After", t)
		http.Error(w, "rate limited", http.StatusTooManyRequests)
	}
}