package middleware

import (
	"net/http"
	"strings"
)

// CSRFProtect returns middleware that validates the Origin header on
// state-changing methods (POST, PUT, PATCH, DELETE). Requests from origins
// not in the trustedOrigins list are rejected with 403. Non-browser clients
// (curl, bots) that send no Origin or Referer header are allowed through,
// since they are not subject to CSRF.
func CSRFProtect(trustedOrigins []string) func(http.Handler) http.Handler {
	originSet := make(map[string]struct{}, len(trustedOrigins))
	for _, o := range trustedOrigins {
		originSet[o] = struct{}{}
	}

	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			switch r.Method {
			case "POST", "PUT", "PATCH", "DELETE":
				origin := r.Header.Get("Origin")
				if origin == "" {
					// Fallback: extract origin from Referer
					referer := r.Header.Get("Referer")
					if referer != "" {
						for _, o := range trustedOrigins {
							if strings.HasPrefix(referer, o) {
								origin = o
								break
							}
						}
					}
				}

				if origin != "" {
					if _, ok := originSet[origin]; !ok {
						http.Error(w, `{"error":"forbidden origin"}`, http.StatusForbidden)
						return
					}
				}
				// If both Origin and Referer are empty, allow through
				// (non-browser client like curl, bots, mobile apps)
			}
			next.ServeHTTP(w, r)
		})
	}
}