package api

import (
	"net/http"
	"strings"
	"time"

	"github.com/jfraeys/fetch_ml/internal/middleware"
)

// wrapWithMiddleware wraps the handler with security middleware
func (s *Server) wrapWithMiddleware(mux *http.ServeMux) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Skip auth for WebSocket and health endpoints
		if r.URL.Path == "/ws" || strings.HasPrefix(r.URL.Path, "/health") {
			mux.ServeHTTP(w, r)
			return
		}

		handler := s.sec.APIKeyAuth(mux)
		handler = s.sec.RateLimit(handler)
		handler = middleware.SecurityHeaders(handler)
		handler = middleware.CORS(s.config.Security.AllowedOrigins)(handler)
		handler = middleware.RequestTimeout(30 * time.Second)(handler)
		handler = middleware.AuditLogger(handler)
		if len(s.config.Security.IPWhitelist) > 0 {
			handler = s.sec.IPWhitelist(s.config.Security.IPWhitelist)(handler)
		}
		handler.ServeHTTP(w, r)
	})
}