package com.zhentao.handler;

import cn.hutool.core.util.IdUtil;
import cn.hutool.json.JSONUtil;
import com.zhentao.constant.ChatConstants;
import com.zhentao.dto.WsMessage;
import com.zhentao.service.ChatMessageService;
import com.zhentao.service.OnlineUserService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 聊天WebSocket处理器
 */
@Component
public class ChatWebSocketHandler implements WebSocketHandler {

    @Autowired
    @Lazy
    private ChatMessageService chatMessageService;

    @Autowired
    private OnlineUserService onlineUserService;

    /**
     * 存储用户ID和WebSocket会话的映射
     * key: userId, value: WebSocketSession
     */
    private static final ConcurrentHashMap<Long, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();

    /**
     * 存储会话ID和用户ID的映射
     * key: sessionId, value: userId
     */
    private static final ConcurrentHashMap<String, Long> SESSION_USERS = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long userId = (Long) session.getAttributes().get("userId");
        
        if (userId != null) {
            // 存储会话
            USER_SESSIONS.put(userId, session);
            SESSION_USERS.put(session.getId(), userId);
            
            // 标记用户在线
            onlineUserService.userOnline(userId, session.getId());
            
            System.out.println("用户 " + userId + " WebSocket连接成功，sessionId: " + session.getId());
            
            // 发送上线通知
            WsMessage onlineMsg = new WsMessage();
            onlineMsg.setType(ChatConstants.WsMessageType.ONLINE);
            onlineMsg.setFromUserId(userId);
            onlineMsg.setTimestamp(System.currentTimeMillis());
            sendMessage(session, onlineMsg);
            
            // 推送离线消息
            chatMessageService.pushOfflineMessages(userId, session);
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        Long userId = SESSION_USERS.get(session.getId());
        
        if (userId == null) {
            System.err.println("未找到会话对应的用户ID: " + session.getId());
            return;
        }
        
        String payload = message.getPayload().toString();
        System.out.println("收到用户 " + userId + " 的消息: " + payload);
        
        try {
            WsMessage wsMessage = JSONUtil.toBean(payload, WsMessage.class);
            wsMessage.setFromUserId(userId);
            
            // 根据消息类型处理
            switch (wsMessage.getType()) {
                case ChatConstants.WsMessageType.PING:
                    handlePing(session, wsMessage);
                    break;
                case ChatConstants.WsMessageType.CHAT:
                    handleChatMessage(session, wsMessage);
                    break;
                case ChatConstants.WsMessageType.ACK:
                    handleAck(session, wsMessage);
                    break;
                case ChatConstants.WsMessageType.RECALL:
                    handleRecall(session, wsMessage);
                    break;
                case ChatConstants.WsMessageType.TYPING:
                    handleTyping(session, wsMessage);
                    break;
                case ChatConstants.WsMessageType.READ_RECEIPT:
                    handleReadReceipt(session, wsMessage);
                    break;
                default:
                    System.err.println("未知的消息类型: " + wsMessage.getType());
            }
        } catch (Exception e) {
            System.err.println("处理消息异常: " + e.getMessage());
            e.printStackTrace();
            
            // 发送错误消息
            WsMessage errorMsg = new WsMessage();
            errorMsg.setType(ChatConstants.WsMessageType.ERROR);
            errorMsg.setErrorMessage("消息处理失败: " + e.getMessage());
            sendMessage(session, errorMsg);
        }
    }

    /**
     * 处理心跳ping消息
     */
    private void handlePing(WebSocketSession session, WsMessage message) throws IOException {
        Long userId = message.getFromUserId();
        
        // 更新用户在线状态
        onlineUserService.updateHeartbeat(userId);
        
        // 回复pong
        WsMessage pong = new WsMessage();
        pong.setType(ChatConstants.WsMessageType.PONG);
        pong.setTimestamp(System.currentTimeMillis());
        sendMessage(session, pong);
    }

    /**
     * 处理聊天消息
     */
    private void handleChatMessage(WebSocketSession session, WsMessage message) {
        // 生成消息ID
        if (message.getMessageId() == null || message.getMessageId().isEmpty()) {
            message.setMessageId(IdUtil.simpleUUID());
        }
        
        // 保存消息并转发
        chatMessageService.saveAndForwardMessage(message, session);
    }

    /**
     * 处理消息确认ACK
     */
    private void handleAck(WebSocketSession session, WsMessage message) {
        chatMessageService.updateMessageStatus(
            message.getMessageId(), 
            ChatConstants.SendStatus.DELIVERED
        );
    }

    /**
     * 处理消息撤回
     */
    private void handleRecall(WebSocketSession session, WsMessage message) {
        chatMessageService.recallMessage(message.getMessageId(), message.getFromUserId(), session);
    }

    /**
     * 处理正在输入状态
     */
    private void handleTyping(WebSocketSession session, WsMessage message) {
        // 直接转发给对方用户
        WebSocketSession targetSession = USER_SESSIONS.get(message.getToUserId());
        if (targetSession != null && targetSession.isOpen()) {
            sendMessage(targetSession, message);
        }
    }

    /**
     * 处理已读回执
     */
    private void handleReadReceipt(WebSocketSession session, WsMessage message) {
        chatMessageService.markMessagesAsRead(
            message.getFromUserId(), 
            message.getToUserId()
        );
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        Long userId = SESSION_USERS.get(session.getId());
        System.err.println("WebSocket传输错误，用户: " + userId + ", 错误: " + exception.getMessage());
        
        if (session.isOpen()) {
            session.close();
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        Long userId = SESSION_USERS.get(session.getId());
        
        if (userId != null) {
            // 移除会话
            USER_SESSIONS.remove(userId);
            SESSION_USERS.remove(session.getId());
            
            // 标记用户离线
            onlineUserService.userOffline(userId, session.getId());
            
            System.out.println("用户 " + userId + " WebSocket连接关闭，原因: " + closeStatus);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 发送消息到WebSocket客户端
     */
    private void sendMessage(WebSocketSession session, WsMessage message) {
        if (session != null && session.isOpen()) {
            try {
                String json = JSONUtil.toJsonStr(message);
                session.sendMessage(new TextMessage(json));
            } catch (IOException e) {
                System.err.println("发送WebSocket消息失败: " + e.getMessage());
            }
        }
    }

    /**
     * 向指定用户发送消息
     */
    public void sendMessageToUser(Long userId, WsMessage message) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        if (session != null) {
            sendMessage(session, message);
        }
    }

    /**
     * 获取用户的WebSocket会话
     */
    public static WebSocketSession getUserSession(Long userId) {
        return USER_SESSIONS.get(userId);
    }

    /**
     * 检查用户是否在线
     */
    public static boolean isUserOnline(Long userId) {
        WebSocketSession session = USER_SESSIONS.get(userId);
        return session != null && session.isOpen();
    }
}


