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(allowedOrigins []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")

			// Block requests with no referer (this blocks direct browser requests)
			if referer == "" {
				http.Error(w, "Forbidden!", http.StatusForbidden)
				return
			}

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

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