package middleware

import (
	"net/http"

	"github.com/gorilla/sessions"
)

// SecurityHeaders adds security and cache-control headers to all responses
func SecureHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Set strict cache control headers
		w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate, max-age=0")
		w.Header().Set("Pragma", "no-cache")
		w.Header().Set("Expires", "0")

		// Add security headers with updated CSP
		w.Header().Set("Content-Security-Policy",
			"default-src 'self'; "+
				"script-src 'self' https://unpkg.com/htmx.org@* 'unsafe-inline'; "+
				"style-src 'self' 'unsafe-inline'; "+
				"img-src 'self' data:; "+
				"connect-src 'self'; "+
				"frame-ancestors 'none'; "+
				"form-action 'self'")
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("X-Frame-Options", "DENY")
		w.Header().Set("X-XSS-Protection", "1; mode=block")

		next.ServeHTTP(w, r)
	})
}

// middleware/csrf.go
func CSRFMiddleware(store sessions.Store) Middleware {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Method == "POST" {
				session, _ := store.Get(r, "auth-session")
				csrfToken := session.Values["csrf_token"].(string)
				formToken := r.FormValue("_csrf")

				if csrfToken != formToken {
					http.Error(w, "Invalid CSRF token", http.StatusForbidden)
					return
				}
			}
			next.ServeHTTP(w, r)
		})
	}
}

// MaxBodySize limits the maximum size of request bodies
// size parameter is in bytes
func MaxBodySize(maxSize int64) Middleware {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Skip restricting GET, HEAD, and OPTIONS requests as they shouldn't have bodies
			if r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodOptions {
				next.ServeHTTP(w, r)
				return
			}

			// Check Content-Length header first for efficiency
			if r.ContentLength > maxSize {
				http.Error(w, "Request body too large", http.StatusRequestEntityTooLarge)
				return
			}

			// If Content-Length is not set or potentially spoofed, use LimitReader
			r.Body = http.MaxBytesReader(w, r.Body, maxSize)

			// Continue to next middleware/handler
			next.ServeHTTP(w, r)
		})
	}
}