backend/internal/handler/ws_handler.go

package handler

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

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

	"with_you/internal/dto"
	apperrors "with_you/internal/errors"
	"with_you/internal/model"
	"with_you/internal/pkg/response"
	"with_you/internal/pkg/ws"
	"with_you/internal/repository"
	"with_you/internal/service"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 4096,
	CheckOrigin: func(r *http.Request) bool {
		origin := r.Header.Get("Origin")
		if origin == "" {
			return true
		}
		return isAllowedWebSocketOrigin(origin)
	},
	HandshakeTimeout: 10 * time.Second,
}

var allowedWebSocketOrigins = []string{
	"https://withyou.littlelan.cn",
	"https://admin.littlelan.cn",
	"https://browser.littlelan.cn",
	"http://localhost:3000",
	"http://localhost:5173",
	"http://127.0.0.1:3000",
	"http://127.0.0.1:5173",
}

func isAllowedWebSocketOrigin(origin string) bool {
	for _, o := range allowedWebSocketOrigins {
		if o == origin {
			return true
		}
	}

	if strings.HasPrefix(origin, "http://localhost") || strings.HasPrefix(origin, "http://127.0.0.1") {
		return true
	}

	if strings.HasPrefix(origin, "https://") && strings.HasSuffix(origin, ".littlelan.cn") {
		return true
	}

	return false
}

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

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

// 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 != "" {
			prefix, rest, found := strings.Cut(authHeader, " ")
			if found && prefix == "Bearer" {
				token = rest
			}
		}
	}

	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

	zap.L().Info("WebSocket connection attempt",
		zap.String("user_id", 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{}),
		PendingAcks: make(map[string]time.Time),
	}

	// 4. 注册客户端
	regErr := h.wsHub.Register(client)
	if regErr != nil {
		zap.L().Warn("WebSocket registration rejected",
			zap.String("user_id", userID),
			zap.Error(regErr),
		)
		conn.WriteMessage(websocket.CloseMessage,
			[]byte(`{"type":"error","payload":{"code":"connection_limit","message":"too many connections"}}`))
		conn.Close()
		return
	}
	defer h.wsHub.Unregister(client)

	zap.L().Info("WebSocket client connected",
		zap.String("user_id", userID),
		zap.Uint64("client_id", clientID),
	)

	// 5. 提示客户端进行 seq 同步
	syncMsg := ws.ResponseMessage{
		EventID: h.wsHub.NextID(),
		Type:    "sync_required",
		TS:      time.Now().UnixMilli(),
	}
	if syncData, err := json.Marshal(syncMsg); err == nil {
		select {
		case <-client.Quit:
		case client.Send <- syncData:
		}
	}

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

	zap.L().Info("WebSocket client disconnected",
		zap.String("user_id", userID),
		zap.Uint64("client_id", clientID),
	)
}

// readPump 读取客户端消息
func (h *WSHandler) readPump(conn *websocket.Conn, client *ws.Client) {
	defer func() {
		zap.L().Info("WebSocket readPump exiting",
			zap.String("user_id", client.UserID),
			zap.Uint64("client_id", client.ID),
		)
		conn.Close()
	}()

	conn.SetReadLimit(256 * 1024) // 256KB (SDP/ICE candidates can be large)
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(string) error {
		zap.L().Debug("WebSocket pong received",
			zap.String("user_id", client.UserID),
		)
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	zap.L().Info("WebSocket readPump started",
		zap.String("user_id", client.UserID),
		zap.Uint64("client_id", client.ID),
	)

	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),
				)
			} else {
				zap.L().Info("WebSocket connection closed",
					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
		}

		zap.L().Debug("WebSocket message received",
			zap.String("user_id", client.UserID),
			zap.String("type", msg.Type),
		)

		// 处理消息
		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 {
				zap.L().Warn("WebSocket write error (NextWriter)",
					zap.String("user_id", client.UserID),
					zap.Uint64("client_id", client.ID),
					zap.Error(err),
				)
				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 {
				zap.L().Warn("WebSocket write error (Close)",
					zap.String("user_id", client.UserID),
					zap.Uint64("client_id", client.ID),
					zap.Error(err),
				)
				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)
	case "ack":
		h.handleAck(client, msg.Payload)
	// 通话信令
	case "call_invite":
		h.handleCallInvite(ctx, client, msg.Payload)
	case "call_answer":
		h.handleCallAnswer(ctx, client, msg.Payload)
	case "call_reject":
		h.handleCallReject(ctx, client, msg.Payload)
	case "call_busy":
		h.handleCallBusy(ctx, client, msg.Payload)
	case "call_sdp":
		h.handleCallSDP(ctx, client, msg.Payload)
	case "call_ice":
		h.handleCallICE(ctx, client, msg.Payload)
	case "call_end":
		h.handleCallEnd(ctx, client, msg.Payload)
	case "call_mute":
		h.handleCallMute(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) {
	if !h.isVerified(ctx, client) {
		return
	}

	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"`
		ClientMsgID    string                `json:"client_msg_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, req.ClientMsgID)
	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,
		ConversationID: req.ConversationID,
		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) {
	if !h.isVerified(ctx, client) {
		return
	}

	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

// handleAck 处理客户端ACK确认
func (h *WSHandler) handleAck(client *ws.Client, payload json.RawMessage) {
	var req struct {
		MessageID string `json:"message_id"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.MessageID == "" {
		return
	}
	h.wsHub.AckMessage(client.UserID, req.MessageID)
}

// isVerified 检查用户是否已通过身份认证
func (h *WSHandler) isVerified(ctx context.Context, client *ws.Client) bool {
	user, err := h.userRepo.GetByID(client.UserID)
	if err != nil {
		h.wsHub.SendError(client, "internal_error", "获取用户信息失败")
		return false
	}
	if user.VerificationStatus != model.VerificationStatusApproved {
		h.wsHub.SendError(client, "VERIFICATION_REQUIRED", "请先完成身份认证")
		return false
	}
	return true
}

// ==================== 通话信令处理 ====================

// handleCallInvite 处理呼叫邀请
func (h *WSHandler) handleCallInvite(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	if !h.isVerified(ctx, client) {
		return
	}

	var req struct {
		CalleeID       string `json:"callee_id"`
		ConversationID string `json:"conversation_id"`
		MediaType      string `json:"call_type"` // voice 或 video
	}
	if err := json.Unmarshal(payload, &req); err != nil {
		h.wsHub.SendError(client, "parse_error", "invalid call_invite format")
		return
	}
	if req.CalleeID == "" || req.ConversationID == "" {
		h.wsHub.SendError(client, "invalid_params", "callee_id and conversation_id are required")
		return
	}

	// 前端传的 call_type 是 voice/video，默认为 voice
	mediaType := req.MediaType
	if mediaType != "video" {
		mediaType = "voice"
	}

	call, _, err := h.callService.Invite(ctx, client.UserID, req.CalleeID, req.ConversationID, mediaType)
	if err != nil {
		zap.L().Warn("Failed to invite call",
			zap.String("caller_id", client.UserID),
			zap.String("callee_id", req.CalleeID),
			zap.Error(err),
		)
		if errors.Is(err, apperrors.ErrCallInProgress) {
			h.wsHub.SendError(client, "call_in_progress", err.Error())
			return
		}
		h.wsHub.SendError(client, "call_invite_failed", err.Error())
		return
	}

	// 返回给呼叫方的响应
	resp := ws.ResponseMessage{
		EventID: h.wsHub.NextID(),
		Type:    "call_invited",
		TS:      time.Now().UnixMilli(),
		Payload: map[string]any{
			"call_id":         call.ID,
			"conversation_id": req.ConversationID,
			"callee_id":       req.CalleeID,
			"lifetime":        service.CallLifetimeMs,
		},
	}
	data, _ := json.Marshal(resp)
	select {
	case <-client.Quit:
	case client.Send <- data:
	}
}

// handleCallAnswer 处理接听
func (h *WSHandler) handleCallAnswer(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		CallID string `json:"call_id"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.CallID == "" {
		h.wsHub.SendError(client, "invalid_params", "call_id is required")
		return
	}

	call, err := h.callService.Accept(ctx, req.CallID, client.UserID)
	if err != nil {
		// === 区分已接听错误 ===
		if errors.Is(err, apperrors.ErrCallAlreadyAnswered) {
			h.wsHub.SendError(client, "call_already_answered", "通话已被其他设备接听")
			return
		}
		h.wsHub.SendError(client, "call_accept_failed", err.Error())
		return
	}

	resp := ws.ResponseMessage{
		EventID: h.wsHub.NextID(),
		Type:    "call_accepted",
		TS:      time.Now().UnixMilli(),
		Payload: map[string]any{
			"call_id":    call.ID,
			"started_at": call.StartedAt,
		},
	}
	data, _ := json.Marshal(resp)
	select {
	case <-client.Quit:
	case client.Send <- data:
	}
}

// handleCallReject 处理拒绝
func (h *WSHandler) handleCallReject(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		CallID string `json:"call_id"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.CallID == "" {
		h.wsHub.SendError(client, "invalid_params", "call_id is required")
		return
	}

	if err := h.callService.Reject(ctx, req.CallID, client.UserID); err != nil {
		h.wsHub.SendError(client, "call_reject_failed", err.Error())
	}
}

// handleCallBusy 处理忙线
func (h *WSHandler) handleCallBusy(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		CallID string `json:"call_id"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.CallID == "" {
		h.wsHub.SendError(client, "invalid_params", "call_id is required")
		return
	}

	if err := h.callService.Busy(ctx, req.CallID, client.UserID); err != nil {
		h.wsHub.SendError(client, "call_busy_failed", err.Error())
	}
}

// handleCallSDP 转发 SDP
func (h *WSHandler) handleCallSDP(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		CallID  string          `json:"call_id"`
		SDPType string          `json:"sdp_type"`
		SDP     json.RawMessage `json:"sdp"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.CallID == "" || req.SDPType == "" {
		h.wsHub.SendError(client, "invalid_params", "call_id and sdp_type are required")
		return
	}

	signalPayload, _ := json.Marshal(map[string]any{
		"sdp_type": req.SDPType,
		"sdp":      json.RawMessage(req.SDP),
	})
	if err := h.callService.RelaySignal(ctx, req.CallID, client.UserID, "call_sdp", signalPayload); err != nil {
		h.wsHub.SendError(client, "call_sdp_failed", err.Error())
	}
}

// handleCallICE 转发 ICE candidate
func (h *WSHandler) handleCallICE(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		CallID    string          `json:"call_id"`
		Candidate json.RawMessage `json:"candidate"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.CallID == "" {
		h.wsHub.SendError(client, "invalid_params", "call_id is required")
		return
	}

	signalPayload, _ := json.Marshal(map[string]any{
		"candidate": json.RawMessage(req.Candidate),
	})
	if err := h.callService.RelaySignal(ctx, req.CallID, client.UserID, "call_ice", signalPayload); err != nil {
		h.wsHub.SendError(client, "call_ice_failed", err.Error())
	}
}

// handleCallEnd 处理结束通话
func (h *WSHandler) handleCallEnd(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		CallID string `json:"call_id"`
		Reason string `json:"reason"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.CallID == "" {
		h.wsHub.SendError(client, "invalid_params", "call_id is required")
		return
	}

	call, err := h.callService.End(ctx, req.CallID, client.UserID, req.Reason)
	if err != nil {
		h.wsHub.SendError(client, "call_end_failed", err.Error())
		return
	}

	resp := ws.ResponseMessage{
		EventID: h.wsHub.NextID(),
		Type:    "call_ended",
		TS:      time.Now().UnixMilli(),
		Payload: map[string]any{
			"call_id":  call.ID,
			"duration": call.Duration,
		},
	}
	data, _ := json.Marshal(resp)
	select {
	case <-client.Quit:
	case client.Send <- data:
	}
}

// handleCallMute 处理静音
func (h *WSHandler) handleCallMute(ctx context.Context, client *ws.Client, payload json.RawMessage) {
	var req struct {
		CallID string `json:"call_id"`
		Muted  bool   `json:"muted"`
	}
	if err := json.Unmarshal(payload, &req); err != nil || req.CallID == "" {
		h.wsHub.SendError(client, "invalid_params", "call_id is required")
		return
	}

	if err := h.callService.SetMuted(ctx, req.CallID, client.UserID, req.Muted); err != nil {
		h.wsHub.SendError(client, "call_mute_failed", err.Error())
	}
}