package main

import (
	"net/http"
	"strings"

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

// RateLimiter middleware limits requests per IP
func RateLimiter(rps float64, burst int) func(http.Handler) http.Handler {
	limiter := rate.NewLimiter(rate.Limit(rps), burst)
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			if !limiter.Allow() {
				http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
				return
			}
			next.ServeHTTP(w, req)
		})
	}
}

func RefererCheck(allowedDomains []string) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			referer := r.Header.Get("Referer")

			// Allow requests with no referer (direct API calls, Postman, etc.)
			if referer == "" {
				next.ServeHTTP(w, r)
				return
			}

			// Check if referer matches allowed domains
			for _, domain := range allowedDomains {
				if strings.HasPrefix(referer, domain) {
					next.ServeHTTP(w, r)
					return
				}
			}

			// Block if referer doesn't match
			http.Error(w, "Forbidden", http.StatusForbidden)
		})
	}
}