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.
//
// If host is non-empty, any Origin whose hostname matches host is allowed.
// additionalOrigins is a list of extra fully-qualified origins to trust
// (e.g. from environment variable DUMPSTER_CSRF_ORIGINS, comma-separated).
func CSRFProtect(host, port string, additionalOrigins []string) func(http.Handler) http.Handler {
	// Build a set of exact-match origins for quick lookup
	originSet := make(map[string]struct{})
	for _, o := range additionalOrigins {
		o = strings.TrimSpace(o)
		if o != "" {
			originSet[o] = struct{}{}
		}
	}

	// Extract hostname from host (strip port if present)
	hostname := host
	if idx := strings.LastIndex(hostname, ":"); idx > 0 {
		hostname = hostname[:idx]
	}

	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 originSet {
							if strings.HasPrefix(referer, o) {
								origin = o
								break
							}
						}
					}
				}

				if origin != "" {
					if !isOriginTrusted(origin, hostname, port, originSet) {
						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)
		})
	}
}

// isOriginTrusted checks if the given Origin header value is trusted.
// An origin is trusted if:
// 1. It's in the explicit originSet, OR
// 2. Its hostname matches the configured host (any scheme/port)
func isOriginTrusted(origin, host, port string, originSet map[string]struct{}) bool {
	if _, ok := originSet[origin]; ok {
		return true
	}

	// Parse the origin to extract hostname and port
	// Origin format: scheme://host[:port]
	withoutScheme := origin
	if idx := strings.Index(withoutScheme, "://"); idx >= 0 {
		withoutScheme = withoutScheme[idx+3:]
	}
	// Remove path if any
	if idx := strings.Index(withoutScheme, "/"); idx >= 0 {
		withoutScheme = withoutScheme[:idx]
	}

	originHost := withoutScheme
	originPort := ""
	if idx := strings.LastIndex(originHost, ":"); idx >= 0 {
		originPort = originHost[idx+1:]
		originHost = originHost[:idx]
	}

	// Strip brackets from IPv6
	originHost = strings.TrimPrefix(originHost, "[")
	originHost = strings.TrimSuffix(originHost, "]")

	if host != "" && originHost == host {
		return true
	}

	// Also trust localhost variants
	if originHost == "localhost" || originHost == "127.0.0.1" || originHost == "::1" {
		if port == "" || originPort == "" || originPort == port {
			return true
		}
	}

	return false
}