package cn.reghao.bnt.web.ws.handler;

import cn.reghao.bnt.common.msg.constant.NodeStatus;
import cn.reghao.bnt.web.devops.machine.service.MachineService;
import cn.reghao.bnt.web.devops.machine.util.Keys;
import cn.reghao.bnt.web.ws.EventCenter;
import cn.reghao.bnt.web.ws.WsSender;
import cn.reghao.jutil.jdk.event.message.EventMessage;
import cn.reghao.jutil.jdk.serializer.JdkSerializer;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.nio.ByteBuffer;

/**
 * @author reghao
 * @date 2023-02-22 18:14:56
 */
@Slf4j
@Component
public class AgentWebSocketHandler implements WebSocketHandler {
    private final WsSender wsSender;
    private final MachineService machineService;
    private final EventCenter eventCenter;

    public AgentWebSocketHandler(WsSender wsSender, EventCenter eventCenter, MachineService machineService) {
        this.wsSender = wsSender;
        this.eventCenter = eventCenter;
        this.machineService = machineService;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) {
        String machineId = getId(webSocketSession);
        log.info("节点 {} 建立连接", machineId);
        wsSender.addSession(machineId, webSocketSession);
        machineService.setAgentStatus(machineId, NodeStatus.Online);
    }

    private String getId(WebSocketSession webSocketSession) {
        return (String) webSocketSession.getAttributes().get(Keys.USER_UUID);
    }

    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) {
        try {
            if (webSocketMessage instanceof TextMessage) {
                log.info("收到来自 {} 的文本消息", webSocketSession.getId());
                TextMessage textMessage = (TextMessage) webSocketMessage;
            } else if (webSocketMessage instanceof BinaryMessage) {
                //log.info("收到来自 {} 的二进制消息", webSocketSession.getId());
                BinaryMessage binaryMessage = (BinaryMessage) webSocketMessage;
                ByteBuffer byteBuffer = binaryMessage.getPayload();
                byte[] bytes = byteBuffer.array();

                Object obj = JdkSerializer.deserialize(bytes);
                if (obj instanceof EventMessage) {
                    EventMessage eventMessage = (EventMessage) obj;
                    eventCenter.dispatch(eventMessage);
                }
            } else if (webSocketMessage instanceof PingMessage) {
                log.info("收到来自 {} 的 PingMessage 消息", webSocketSession.getId());
                webSocketSession.sendMessage(new PongMessage());
            } else if (webSocketMessage instanceof PongMessage) {
                log.info("收到来自 {} 的 PongMessage 消息", webSocketSession.getId());
            } else {
                log.error("未知类型的 WebSocket 消息");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) {
        log.error("WebSocket 数据传输错误");
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) {
        String machineId = getId(webSocketSession);
        machineService.setAgentStatus(machineId, NodeStatus.Offline);
        wsSender.removeSession(machineId);
        log.info("节点 {} 断开连接", machineId);
    }

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