dating/src/main/java/com/bao/dating/controller/websocket/ChatWebSocketHandler.java

package com.bao.dating.controller.websocket;

import com.bao.dating.message.WsMessage;
import com.bao.dating.pojo.dto.ChatRecallDTO;
import com.bao.dating.pojo.dto.ChatRecordSendDTO;
import com.bao.dating.pojo.vo.ChatRecordsVO;
import com.bao.dating.service.ChatService;
import com.bao.dating.session.WsSessionManager;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

/**
 * WebSocket处理器类
 * @author lenovo
 */
@Slf4j
@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {

    @Autowired
    private ObjectMapper objectMapper;

    @Autowired
    private ChatService chatService;

    @Autowired
    private WsSessionManager sessionManager;


    /**
     * 用户建立连接（上线）
     * @param session WebSocketSession
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        Long userId = (Long) session.getAttributes().get("userId");
        sessionManager.addSession(userId, session);
        log.info("用户 " + userId + " 已上线");
    }

    /**
     * 接收并处理消息
     * @param session WebSocketSession
     * @param message 消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception{
        // 获取当前用户ID
        Long senderUserId = (Long) session.getAttributes().get("userId");
        if (senderUserId == null) {
            log.error("WebSocket session 中未找到 userId");
            return;
        }

        JsonNode node = objectMapper.readTree(message.getPayload());
        String type = node.get("type").asText();
        // 根据消息类型解析消息
        WsMessage wsMessage = objectMapper.readValue(message.getPayload(), WsMessage.class);
        
        // 先获取消息类型，再根据类型进行相应处理和转换
        if ("chat".equals(type)) {
            // 处理私聊消息
            WsMessage<ChatRecordSendDTO> chatWsMessage = 
                objectMapper.convertValue(node, new TypeReference<WsMessage<ChatRecordSendDTO>>(){});
            handlePrivateChat(session, senderUserId, chatWsMessage.getData());
        } else if ("recall".equals(type)) {
            // 处理撤回消息
            WsMessage<ChatRecallDTO> recallWsMessage = 
                objectMapper.convertValue(node, new TypeReference<WsMessage<ChatRecallDTO>>(){});
            handleRecallMessage(session, senderUserId, recallWsMessage.getData());
        }
    }

    /**
     * 私聊处理
     */
    private void handlePrivateChat(WebSocketSession session, Long senderUserId, ChatRecordSendDTO dto) throws Exception {

        // 1. 消息入库 + 会话更新
        ChatRecordsVO chatRecordsVO = chatService.createSession(senderUserId, dto);
        if (chatRecordsVO == null){
            WsMessage<String> errorMsg = new WsMessage<>();
            errorMsg.setType("error");
            errorMsg.setData("会话已删除，无法发送消息");
            // 返回错误信息
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMsg)));
            return;
        }
        // 2. 推送给接收方
        WebSocketSession receiverSession =
                sessionManager.getSession(dto.getReceiverUserId());

        if (receiverSession != null && receiverSession.isOpen()) {
            WsMessage<ChatRecordsVO> pushMsg = new WsMessage<>();
            pushMsg.setType("chat");
            pushMsg.setData(chatRecordsVO);

            receiverSession.sendMessage(
                    new TextMessage(objectMapper.writeValueAsString(pushMsg))
            );
        }
    }

    /**
     * 消息撤回处理
     */
    private void handleRecallMessage(WebSocketSession session, Long senderUserId, Object data) throws Exception {

        // 转 DTO
        ChatRecallDTO dto = objectMapper.convertValue(data, ChatRecallDTO.class);

        // 撤回逻辑
        boolean success = chatService.recallMessage(senderUserId, dto.getChatId());
        // 如果返回false，说明消息撤回失败
        if (!success) {
            WsMessage<String> errorMsg = new WsMessage<>();
            errorMsg.setType("error");
            errorMsg.setData("撤回失败，消息可能无法撤回或不存在");
            // 返回错误信息
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(errorMsg)));
            return;
        }

        // 创建撤回通知消息
        WsMessage<ChatRecallDTO> pushMsg = new WsMessage<>();
        pushMsg.setType("recall");
        pushMsg.setData(dto);

        // 通知自己
        if (session.isOpen()) {
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(pushMsg)));
        }

        // 通知对方
        WebSocketSession receiverSession =
                sessionManager.getSession(dto.getReceiverUserId());

        if (receiverSession != null && receiverSession.isOpen()) {
            receiverSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(pushMsg))
            );
        }
    }


    /**
     * 用户断开连接（下线）
     * @param session WebSocketSession
     * @param status 断开原因
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status){
        // 下线处理
        Long userId = (Long) session.getAttributes().get("userId");
        sessionManager.removeSession(userId);
        log.info("用户 " + userId + " 已下线");
    }
}