package com.ruoyi.dj.websocket.config;

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.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

/**
 * 自定义WebSocket处理类
 */
@Slf4j
@Component
public class CustomWebSocketHandler extends TextWebSocketHandler {
    /**
     * 存储所有WebSocket会话，以确保线程安全
     */
    private Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    /**
     * Redis中存储会话ID与客户端选择的gateway的映射关系的前缀
     */
    private static final String SESSION_GATEWAY_PREFIX = "websocket:session:gateway:";

    private ObjectMapper objectMapper = new ObjectMapper();

    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 在连接建立后调用，将会话添加到会话映射中
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.put(session.getId(), session);
        log.info("会话已建立，sessionId={}", session.getId());
    }

    /**
     * 处理接收到客户端的文本消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        JsonNode jsonNode = objectMapper.readTree(payload);

        String clientID = jsonNode.has("clientID") ? jsonNode.get("clientID").asText() : null;
        String gateway = jsonNode.has("gateway") ? jsonNode.get("gateway").asText() : null;

        if (clientID != null && gateway != null) {
            // 使用Redis存储会话ID与gateway的映射关系
            String key = SESSION_GATEWAY_PREFIX + session.getId();
            redisTemplate.opsForValue().set(key, gateway);
            log.info("会话ID={} 已选择网关={}", session.getId(), gateway);
        }
    }

    /**
     * 向所有选择了特定 Gateway 的客户端发送消息
     */
    public void sendMessageToAll(String message) {
        try {
            JsonNode jsonNode = objectMapper.readTree(message);
            String gateway = jsonNode.has("gateway") ? jsonNode.get("gateway").asText() : null;
            if (gateway == null) {
                log.warn("消息中缺少gateway字段，无法发送");
                return;
            }

            // 创建文本消息对象
            TextMessage textMessage = new TextMessage(message);

            // 遍历所有会话，尝试发送消息给选择了对应gateway的会话
            sessions.values().forEach(session -> {
                try {
                    if (session.isOpen()) {
                        String key = SESSION_GATEWAY_PREFIX + session.getId();
                        String sessionGateway = redisTemplate.opsForValue().get(key);
                        if (gateway.equals(sessionGateway)) {
                            session.sendMessage(textMessage);
                        }
                    }
                } catch (Exception e) {
                    log.error("发送消息给sessionId={}时发生异常", session.getId(), e);
                }
            });
        } catch (Exception e) {
            log.error("解析消息时发生异常", e);
        }
    }

    /**
     * 在连接关闭后调用，移除会话及其映射关系
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        sessions.remove(session.getId());
        // 从Redis中删除会话与gateway的映射
        String key = SESSION_GATEWAY_PREFIX + session.getId();
        redisTemplate.delete(key);
        log.info("会话已关闭，sessionId={}", session.getId());
    }

    /**
     * 处理传输错误
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        sessions.remove(session.getId());
        // 从Redis中删除会话与gateway的映射
        String key = SESSION_GATEWAY_PREFIX + session.getId();
        redisTemplate.delete(key);
        log.error("会话传输错误，sessionId={}", session.getId(), exception);
    }
}
