tailly_back_v2/internal/http/handlers/chat.go

package handlers

import (
	"context"
	"log"
	"net/http"
	"tailly_back_v2/internal/domain"
	"tailly_back_v2/internal/service"
	"tailly_back_v2/internal/ws"
	"tailly_back_v2/pkg/auth"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
	Subprotocols: []string{"graphql-transport-ws"},
}

type ChatHandler struct {
	chatService service.ChatService
	hub         *ws.Hub
	tokenAuth   *auth.TokenAuth
}

func NewChatHandler(chatService service.ChatService, hub *ws.Hub, tokenAuth *auth.TokenAuth) *ChatHandler {
	return &ChatHandler{
		chatService: chatService,
		hub:         hub,
		tokenAuth:   tokenAuth,
	}
}

func (h *ChatHandler) HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	log.Printf("Incoming WebSocket headers: %+v", r.Header)
	log.Printf("Cookies: %+v", r.Cookies())

	// Убираем все проверки токена здесь, так как будем проверять его в первом сообщении
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		http.Error(w, "Could not upgrade to WebSocket", http.StatusBadRequest)
		return
	}

	client := &ws.Client{
		UserID: 0, // Пока не авторизован
		Send:   make(chan *domain.Message, 256),
	}

	h.hub.RegisterClient(client)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go h.readPump(ctx, conn, client)
	go h.writePump(ctx, conn, client)

	<-ctx.Done()
}

func (h *ChatHandler) readPump(ctx context.Context, conn *websocket.Conn, client *ws.Client) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	defer conn.Close()

	// 1. Ожидаем первое сообщение с токеном
	var authMsg struct {
		Type  string `json:"type"`
		Token string `json:"token"`
	}

	if err := conn.ReadJSON(&authMsg); err != nil {
		log.Printf("Failed to read auth message: %v", err)
		return
	}

	if authMsg.Type != "auth" || authMsg.Token == "" {
		log.Println("First message must be auth with token")
		conn.WriteJSON(map[string]string{"type": "error", "message": "Auth required"})
		return
	}

	// 2. Валидируем токен
	userID, err := h.tokenAuth.ValidateAccessToken(authMsg.Token)
	if err != nil {
		log.Printf("Token validation error: %v", err)
		conn.WriteJSON(map[string]string{"type": "error", "message": "Invalid token"})
		return
	}

	client.UserID = userID
	log.Printf("WebSocket authenticated, userID=%d", userID)

	// 3. Отправляем подтверждение авторизации
	if err := conn.WriteJSON(map[string]string{"type": "auth_success"}); err != nil {
		log.Printf("Failed to send auth confirmation: %v", err)
		return
	}

	// 4. Основной цикл обработки сообщений
	for {
		select {
		case <-ticker.C:
			if err := conn.WriteJSON(map[string]string{"type": "ping"}); err != nil {
				log.Printf("Ping error: %v", err)
				return
			}
		case <-ctx.Done():
			return
		default:
			var msg struct {
				Type    string `json:"type"`
				Payload struct {
					ReceiverID int    `json:"receiverId"`
					Content    string `json:"content"`
				} `json:"payload"`
			}

			if err := conn.ReadJSON(&msg); err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					log.Printf("WebSocket error: %v", err)
				}
				return
			}

			switch msg.Type {
			case "pong":
				continue
			case "ping":
				conn.WriteJSON(map[string]string{"type": "pong"})
				continue
			case "message":
				if client.UserID == 0 {
					conn.WriteJSON(map[string]string{"type": "error", "message": "Not authenticated"})
					continue
				}

				if msg.Payload.ReceiverID == 0 {
					conn.WriteJSON(map[string]interface{}{
						"type":    "error",
						"message": "Invalid receiver ID",
					})
					continue
				}

				chat, err := h.chatService.GetOrCreateChat(ctx, client.UserID, msg.Payload.ReceiverID)
				if err != nil {
					log.Printf("Chat error: %v", err)
					conn.WriteJSON(map[string]interface{}{
						"type":    "error",
						"message": "Chat error",
						"details": err.Error(),
					})
					continue
				}

				message, err := h.chatService.SendMessage(
					ctx,
					client.UserID,
					chat.ID,
					msg.Payload.Content,
				)
				if err != nil {
					log.Printf("Message send error: %v", err)
					continue
				}

				h.hub.Broadcast(message)
			default:
				conn.WriteJSON(map[string]string{"type": "error", "message": "Unknown message type"})
			}
		}
	}
}

func (h *ChatHandler) writePump(ctx context.Context, conn *websocket.Conn, client *ws.Client) {
	defer conn.Close()

	for {
		select {
		case <-ctx.Done():
			return
		case message, ok := <-client.Send:
			if !ok {
				// Канал закрыт
				conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			if err := conn.WriteJSON(message); err != nil {
				log.Printf("WebSocket write error: %v", err)
				return
			}
		}
	}
}