package middleware

import (
	"context"
	"net/http"
	"strings"
	"tailly_back_v2/pkg/auth"
	"time"
)

func WSAuthMiddleware(tokenAuth *auth.TokenAuth) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.Header.Get("Upgrade") == "websocket" {
				ctx, cancel := context.WithTimeout(r.Context(), 30*time.Minute)
				defer cancel()

				token := extractTokenFromWSRequest(r) // ← Исправленная функция
				if token != "" {
					if userID, err := tokenAuth.ValidateAccessToken(token); err == nil {
						ctx = context.WithValue(ctx, userIDKey, userID)
						r = r.WithContext(ctx)
					} else {
						// Отклоняем соединение при невалидном токене
						http.Error(w, "Invalid token", http.StatusUnauthorized)
						return
					}
				} else {
					// Отклоняем соединение без токена
					http.Error(w, "Authentication required", http.StatusUnauthorized)
					return
				}
			}
			next.ServeHTTP(w, r)
		})
	}
}

func extractTokenFromWSRequest(r *http.Request) string {
	// 1. Проверяем заголовок Authorization (WebSocket использует его)
	if authHeader := r.Header.Get("Authorization"); authHeader != "" {
		if strings.HasPrefix(authHeader, "Bearer ") {
			return strings.TrimPrefix(authHeader, "Bearer ")
		}
		return authHeader // На случай если без Bearer prefix
	}

	// 2. Проверяем query parameter (альтернативный способ)
	if token := r.URL.Query().Get("token"); token != "" {
		return token
	}

	// 3. Проверяем куки (последний вариант)
	if cookie, err := r.Cookie("accessToken"); err == nil {
		return cookie.Value
	}

	return ""
}