package com.ruoyi.framework.socket;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;

import com.alibaba.fastjson2.JSON;
import com.ruoyi.common.core.domain.entity.SysDept;
import com.ruoyi.common.core.domain.model.LoginUser;
import com.ruoyi.framework.web.domain.Message;
import com.ruoyi.framework.web.service.TokenService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class MyWebSocketHandler extends TextWebSocketHandler {
    private static final ObjectMapper objectMapper = new ObjectMapper();

    // 存储已连接的用户
    private static final ConcurrentHashMap<Long, WebSocketSession> clients = new ConcurrentHashMap<>();
    @Autowired
    TokenService tokenService;
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String token = getToken(session);
        LoginUser loginUser = tokenService.getLoginUser(token);
        if (loginUser != null) {
            clients.put(loginUser.getUserId(), session);
            sendToUser(loginUser.getUserId(), new Message("Connection Successful"));
        }else{
            session.close(CloseStatus.SERVER_ERROR);
        }


    }


    // **解析 Token**
    private String getToken(WebSocketSession session) {
        String query = Objects.requireNonNull(session.getUri()).getQuery();
        if (query == null || !query.contains("token=")) return null;
        return query.split("token=")[1];
    }

    /**
     * 获取用户
     * @param session
     * @return
     */
    private Long getUserId(WebSocketSession session) {
        for (Map.Entry<Long, WebSocketSession> entry :  clients.entrySet()) {
            if (entry.getValue().equals(session)) return entry.getKey();
        }
        return null;
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        if (message.getPayload().equals("{\"type\":\"ping\"}")) {
            sendToUser(getUserId(session), new Message("pong"));
        }

        super.handleTextMessage(session, message);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        clients.values().removeIf(c -> c == session);
    }


    /**
     * 发送给指定用户
     * @param userId 用户ID
     * @param message 消息内容
     * @throws IOException
     */
    private void sendToUser(Long userId, Message message) throws IOException {
        WebSocketSession webSocketSession = clients.get(userId);
        if (webSocketSession != null && webSocketSession.isOpen()) {
            webSocketSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
        }
    }

    /**
     * 发送给所有人
     * @param message 消息内容
     * @throws IOException
     */
    private void sendToAll(Message message) throws IOException {
        for (WebSocketSession webSocketSession : clients.values()) {
            if (webSocketSession.isOpen()) {
                webSocketSession.sendMessage(new TextMessage(objectMapper.writeValueAsString(message)));
            }
        }
    }

}