package com.fizz.business.service.client;

import cn.hutool.json.JSONUtil;
import com.fizz.business.constants.enums.WsTypeEnum;
import com.fizz.business.dto.DeviceFieldTrendSubscribeReq;
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.*;
import java.util.concurrent.ConcurrentHashMap;

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

    public static DeviceFieldTrendSessionState getFieldTrendState(String sessionId) {
        return fieldTrendStates.get(sessionId);
    }

    public static Map<String, DeviceFieldTrendSessionState> getAllFieldTrendStates() {
        return fieldTrendStates;
    }

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

    /**
     * 懒加载订阅状态：只用于 device_field_trend
     * sessionId -> state
     */
    private static final ConcurrentHashMap<String, DeviceFieldTrendSessionState> fieldTrendStates = 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);

            // device_field_trend 默认启用懒加载：不主动推送，等前端发送订阅消息
            if (WsTypeEnum.device_field_trend.name().equals(type)) {
                fieldTrendStates.put(sid, DeviceFieldTrendSessionState.emptyLazyOn());
            }

            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"));
            return;
        }

        // device_field_trend 订阅（懒加载）
        String type = (String) session.getAttributes().get("type");
        if (WsTypeEnum.device_field_trend.name().equals(type)) {
            try {
                DeviceFieldTrendSubscribeReq req = JSONUtil.toBean(payload, DeviceFieldTrendSubscribeReq.class);
                DeviceFieldTrendSessionState state = fieldTrendStates.get(session.getId());
                if (state == null) {
                    state = DeviceFieldTrendSessionState.emptyLazyOn();
                    fieldTrendStates.put(session.getId(), state);
                }
                if (req.getLazy() != null) {
                    state.setLazy(req.getLazy());
                }
                if (req.getDeviceCodes() != null) {
                    state.setDeviceCodes(new HashSet<>(req.getDeviceCodes()));
                }
                if (req.getFieldNames() != null) {
                    state.setFieldNames(new HashSet<>(req.getFieldNames()));
                }
                log.info("[websocket]device_field_trend订阅更新 sid={}, lazy={}, deviceCodes.size={}, fieldNames.size={}",
                        session.getId(), state.isLazy(),
                        state.getDeviceCodes() == null ? 0 : state.getDeviceCodes().size(),
                        state.getFieldNames() == null ? 0 : state.getFieldNames().size());
            } catch (Exception e) {
                log.error("[websocket]解析device_field_trend订阅失败 sid={}, payload={}", session.getId(), payload, e);
            }
        }
    }

    /**
     * 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();
                }
                // 如果是 field_trend 连接，清除订阅状态
                if (WsTypeEnum.device_field_trend.name().equals(type)) {
                    fieldTrendStates.remove(sid);
                }
            }
            log.info("[websocket]连接关闭:{}", type + "-" + sid);
        } catch (Exception e) {
            log.error("[websocket]连接关闭错误:{}", e.getMessage(), e);
        }

    }
}