package com.fizz.business.service.client;

import lombok.extern.slf4j.Slf4j;
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.AbstractWebSocketHandler;

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

/**
 * @author chenhao
 * @date 2023/05/19
 */
@Slf4j
@Component
public class TrackWsHandler extends AbstractWebSocketHandler {

    private static final ConcurrentHashMap<String, Map<String, WebSocketSession>> clients = new ConcurrentHashMap<>();

    /**
     * socket 建立成功事件
     *
     * @param session
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        try {
            String sid = session.getId();
            String type = (String) session.getAttributes().get("type");
            Map<String, WebSocketSession> sessionMap = clients.get(type);
            if (Objects.isNull(sessionMap)) {
                sessionMap = new ConcurrentHashMap<>();
            }
            sessionMap.put(sid, session);
            // 添加 session
            clients.put(type, sessionMap);
            log.info("[websocket]建立连接:{}", type + sid);
        } catch (Exception e) {
            log.error("[websocket]建立连接错误:{}", e.getMessage(), e);
        }
    }

    /**
     * 接收消息事件
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        log.info("[websocket]收到客户端[{}]的消息:{}", session.getAttributes().get("type") + "-" + session.getId(), payload);
        // 心跳
        if (payload.equals("ping")) {
            session.sendMessage(new TextMessage("pong"));
        }
    }

    /**
     * socket 断开连接时
     *
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        onClose(session);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("[websocket]连接异常:{}", session.getAttributes().get("type") + "-" + session.getId(), exception);
        onClose(session);
    }

    public ConcurrentHashMap<String, Map<String, WebSocketSession>> getClients() {
        return clients;
    }

    private void onClose(WebSocketSession session) {
        try {
            String sid = session.getId();
            String type = (String) session.getAttributes().get("type");
            if (clients.containsKey(type) && clients.get(type).containsKey(sid)) {
                WebSocketSession webSocketSession = clients.get(type).remove(sid);
                if (webSocketSession != null) {
                    webSocketSession.close();
                }
            }
            log.info("[websocket]连接关闭:{}", type + "-" + sid);
        } catch (Exception e) {
            log.error("[websocket]连接关闭错误:{}", e.getMessage(), e);
        }

    }
}