backend/internal/handler/ws_handler.go

package handler

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"go.uber.org/zap"

	"carrot_bbs/internal/dto"
	"carrot_bbs/internal/model"
	"carrot_bbs/internal/pkg/response"
	"carrot_bbs/internal/pkg/ws"
	"carrot_bbs/internal/service"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 4096,
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源，实际生产环境应该限制
	},
	HandshakeTimeout: 10 * time.Second,
}

// WSHandler WebSocket处理器
type WSHandler struct {
	wsHub        *ws.Hub
	chatService  service.ChatService
	groupService service.GroupService
	jwtService   *service.JWTService
	clientSeq    uint64
}

// NewWSHandler 创建WebSocket处理器
func NewWSHandler(
	wsHub *ws.Hub,
	chatService service.ChatService,
	groupService service.GroupService,
	jwtService *service.JWTService,
) *WSHandler {
	return &WSHandler{
		wsHub:        wsHub,
		chatService:  chatService,
		groupService: groupService,
		jwtService:   jwtService,
	}
}

// HandleWebSocket WebSocket连接处理
// GET /api/v1/realtime/ws?token=xxx
func (h *WSHandler) HandleWebSocket(c *gin.Context) {
	// 1. 认证 - 支持Query参数和Header两种方式
	token := c.Query("token")
	if token == "" {
		authHeader := c.GetHeader("Authorization")
		if authHeader != "" {
			parts := strings.SplitN(authHeader, " ", 2)
			if len(parts) == 2 && parts[0] == "Bearer" {
				token = parts[1]
			}
		}
	}

	if token == "" {
		response.Unauthorized(c, "token is required")
		return
	}

	claims, err := h.jwtService.ParseToken(token)
	if err != nil {
		response.Unauthorized(c, "invalid token")
		return
	}

	userID := claims.UserID

	// 2. 升级HTTP连接为WebSocket
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		zap.L().Error("Failed to upgrade WebSocket connection",
			zap.String("user_id", userID),
			zap.Error(err),
		)
		return
	}

	// 3. 创建客户端
	clientID := atomic.AddUint64(&h.clientSeq, 1)
	client := &ws.Client{
		ID:     clientID,
		UserID: userID,
		Send:   make(chan []byte, defaultUserBufferSize),
		Quit:   make(chan struct{}),
	}

	// 4. 注册客户端
	replayEvents := h.wsHub.Register(client)
	defer h.wsHub.Unregister(client)

	// 5. 发送历史回放消息
	go func() {
		for _, ev := range replayEvents {
			msg := ws.ResponseMessage{
				EventID: ev.ID,
				Type:    ev.Type,
				TS:      ev.TS,
				Payload: ev.Payload,
			}
			data, _ := json.Marshal(msg)
			select {
			case <-client.Quit:
				return
			case client.Send <- data:
			}
		}
	}()

	// 6. 启动读写goroutine
	go h.writePump(conn, client)
	h.readPump(conn, client)
}

// readPump 读取客户端消息
func (h *WSHandler) readPump(conn *websocket.Conn, client *ws.Client) {
	defer func() {
		conn.Close()
		close(client.Quit)
	}()

	conn.SetReadLimit(64 * 1024) // 64KB
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				zap.L().Debug("WebSocket read error",
					zap.String("user_id", client.UserID),
					zap.Error(err),
				)
			}
			break
		}

		// 解析消息
		var msg ws.Message
		if err := json.Unmarshal(message, &msg); err != nil {
			h.wsHub.SendError(client, "parse_error", "invalid message format")
			continue
		}

		// 处理消息
		h.handleMessage(client, &msg)
	}
}

// writePump 向客户端发送消息
func (h *WSHandler) writePump(conn *websocket.Conn, client *ws.Client) {
	ticker := time.NewTicker(30 * time.Second)
	defer func() {
		ticker.Stop()
		conn.Close()
	}()

	for {
		select {
		case <-client.Quit:
			conn.WriteMessage(websocket.CloseMessage, []byte{})
			return
		case message, ok := <-client.Send:
			conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if !ok {
				conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			// 批量发送缓冲区中的消息
			n := len(client.Send)
			for i := 0; i < n; i++ {
				w.Write([]byte{'\n'})
				w.Write(<-client.Send)
			}

			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// handleMessage 处理客户端消息
func (h *WSHandler) handleMessage(client *ws.Client, msg *ws.Message) {
	ctx := context.Background()

	switch msg.Type {
	case "ping":
		h.handlePing(client)
	case "chat":
		h.handleChat(ctx, client, msg.Payload)
	case "typing":
		h.handleTyping(ctx, client, msg.Payload)
	case "read":
		h.handleRead(ctx, client, msg.Payload)
	case "recall":
		h.handleRecall(ctx, client, msg.Payload)
	default:
		h.wsHub.SendError(client, "unknown_type", fmt.Sprintf("unknown message type: %s", msg.Type))
	}
}

// handlePing 处理ping消息
func (h *WSHandler) handlePing(client *ws.Client) {
	pong := ws.ResponseMessage{
		EventID: h.wsHub.NextID(),
		Type:    "pong",
		TS:      time.Now().UnixMilli(),
	}
	data, _ := json.Marshal(pong)
	select {
	case <-client.Quit:
	case client.Send <- data:
	}
}

// handleChat 处理聊天消息
func (h *WSHandler) handleChat(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		ConversationID string                `json:"conversation_id"`
		DetailType     string                `json:"detail_type"`
		Segments       model.MessageSegments `json:"segments"`
		ReplyToID      *string               `json:"reply_to_id,omitempty"`
	}

	if err := json.Unmarshal(payload, &req); err != nil {
		h.wsHub.SendError(client, "parse_error", "invalid chat message format")
		return
	}

	if req.ConversationID == "" {
		h.wsHub.SendError(client, "invalid_params", "conversation_id is required")
		return
	}

	if len(req.Segments) == 0 {
		h.wsHub.SendError(client, "invalid_params", "segments is required")
		return
	}

	// 发送消息
	message, err := h.chatService.SendMessage(ctx, client.UserID, req.ConversationID, req.Segments, req.ReplyToID)
	if err != nil {
		h.wsHub.SendError(client, "send_failed", err.Error())
		return
	}

	// 发送成功响应
	detailType := "private"
	if conv, _ := h.chatService.GetConversationByID(ctx, req.ConversationID, client.UserID); conv != nil && conv.Type == model.ConversationTypeGroup {
		detailType = "group"
	}

	resp := dto.WSEventResponse{
		ID:         message.ID,
		Time:       message.CreatedAt.UnixMilli(),
		Type:       "message",
		DetailType: detailType,
		Seq:        strconv.FormatInt(message.Seq, 10),
		Segments:   req.Segments,
		SenderID:   client.UserID,
	}

	msg := ws.ResponseMessage{
		EventID: h.wsHub.NextID(),
		Type:    "message_sent",
		TS:      time.Now().UnixMilli(),
		Payload: resp,
	}
	data, _ := json.Marshal(msg)
	select {
	case <-client.Quit:
	case client.Send <- data:
	}
}

// handleTyping 处理输入状态
func (h *WSHandler) handleTyping(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		ConversationID string `json:"conversation_id"`
		IsTyping       bool   `json:"is_typing"`
	}

	if err := json.Unmarshal(payload, &req); err != nil {
		return
	}

	if req.ConversationID == "" {
		return
	}

	// 输入状态通常不需要持久化，直接通过ChatService的SendTyping处理
	// 这里简化处理，直接推送
	h.chatService.SendTyping(ctx, client.UserID, req.ConversationID)
}

// handleRead 处理已读标记
func (h *WSHandler) handleRead(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		ConversationID string `json:"conversation_id"`
		Seq            int64  `json:"seq"`
	}

	if err := json.Unmarshal(payload, &req); err != nil {
		return
	}

	if req.ConversationID == "" || req.Seq <= 0 {
		return
	}

	err := h.chatService.MarkAsRead(ctx, req.ConversationID, client.UserID, req.Seq)
	if err != nil {
		zap.L().Debug("Failed to mark as read via WebSocket",
			zap.String("user_id", client.UserID),
			zap.String("conversation_id", req.ConversationID),
			zap.Error(err),
		)
	}
}

// handleRecall 处理消息撤回
func (h *WSHandler) handleRecall(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		MessageID string `json:"message_id"`
	}

	if err := json.Unmarshal(payload, &req); err != nil {
		h.wsHub.SendError(client, "parse_error", "invalid recall message format")
		return
	}

	if req.MessageID == "" {
		h.wsHub.SendError(client, "invalid_params", "message_id is required")
		return
	}

	err := h.chatService.RecallMessage(ctx, req.MessageID, client.UserID)
	if err != nil {
		h.wsHub.SendError(client, "recall_failed", err.Error())
		return
	}

	// 发送撤回成功响应
	resp := ws.ResponseMessage{
		EventID: h.wsHub.NextID(),
		Type:    "message_recalled",
		TS:      time.Now().UnixMilli(),
		Payload: map[string]string{"message_id": req.MessageID},
	}
	data, _ := json.Marshal(resp)
	select {
	case <-client.Quit:
	case client.Send <- data:
	}
}

const defaultUserBufferSize = 128