package com.shzg.project.worn.websocket.manager;

import com.shzg.project.worn.websocket.config.WebSocketProperties;
import com.shzg.project.worn.websocket.domain.WsUserInfo;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.Session;
import java.io.IOException;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Slf4j
@Component
public class WebSocketSessionManager {

    /**
     * sessionId -> 用户信息
     */
    private final ConcurrentHashMap<String, WsUserInfo> sessionUserMap = new ConcurrentHashMap<>();

    /**
     * deptId -> session集合
     */
    private final ConcurrentHashMap<Long, CopyOnWriteArraySet<Session>> deptSessionMap = new ConcurrentHashMap<>();

    /**
     * userId -> session集合
     */
    private final ConcurrentHashMap<Long, Set<Session>> userSessionMap = new ConcurrentHashMap<>();

    /**
     * IP -> session集合
     */
    private final ConcurrentHashMap<String, Set<Session>> ipSessionMap = new ConcurrentHashMap<>();

    /**
     * sessionId -> IP
     */
    private final ConcurrentHashMap<String, String> sessionIpMap = new ConcurrentHashMap<>();

    @Autowired
    private WebSocketProperties properties;

    /**
     * 注册连接
     */
    public boolean register(Session session, WsUserInfo userInfo) {

        String sessionId = session.getId();
        Long userId = userInfo.getUserId();

        String ip = session.getUserProperties().get("ip") != null
                ? session.getUserProperties().get("ip").toString()
                : "unknown";

        // ===== 限流 =====

        if (sessionUserMap.size() >= properties.getMaxConnections().getGlobal()) {
            log.warn("[WebSocket] 全局连接数超限");
            return false;
        }

        Set<Session> userSessions =
                userSessionMap.getOrDefault(userId, ConcurrentHashMap.newKeySet());

        if (userSessions.size() >= properties.getMaxConnections().getUser()) {
            log.warn("[WebSocket] 用户连接数超限 userId={}", userId);
            return false;
        }

        Set<Session> ipSessions =
                ipSessionMap.getOrDefault(ip, ConcurrentHashMap.newKeySet());

        if (ipSessions.size() >= properties.getMaxConnections().getIp()) {
            log.warn("[WebSocket] IP连接数超限 ip={}", ip);
            return false;
        }

        // ===== 写入 =====

        sessionUserMap.put(sessionId, userInfo);
        sessionIpMap.put(sessionId, ip);

        userSessions.add(session);
        userSessionMap.put(userId, userSessions);

        ipSessions.add(session);
        ipSessionMap.put(ip, ipSessions);

        for (Long deptId : userInfo.getDeptIds()) {
            deptSessionMap
                    .computeIfAbsent(deptId, k -> new CopyOnWriteArraySet<>())
                    .add(session);
        }

        log.info("[WebSocket] 连接成功 sessionId={}, userId={}, ip={}, 当前连接数={}",
                sessionId, userId, ip, sessionUserMap.size());

        return true;
    }

    /**
     * 移除连接
     */
    public void remove(Session session) {

        if (session == null) return;

        String sessionId = session.getId();

        WsUserInfo userInfo = sessionUserMap.remove(sessionId);
        String ip = sessionIpMap.remove(sessionId);

        if (userInfo != null) {
            Long userId = userInfo.getUserId();

            Set<Session> userSessions = userSessionMap.get(userId);
            if (userSessions != null) {
                userSessions.remove(session);
                if (userSessions.isEmpty()) {
                    userSessionMap.remove(userId);
                }
            }

            for (Long deptId : userInfo.getDeptIds()) {
                Set<Session> sessions = deptSessionMap.get(deptId);
                if (sessions != null) {
                    sessions.remove(session);
                    if (sessions.isEmpty()) {
                        deptSessionMap.remove(deptId);
                    }
                }
            }
        }

        if (ip != null) {
            Set<Session> ipSessions = ipSessionMap.get(ip);
            if (ipSessions != null) {
                ipSessions.remove(session);
                if (ipSessions.isEmpty()) {
                    ipSessionMap.remove(ip);
                }
            }
        }

        log.info("[WebSocket] 连接移除 sessionId={}, 当前连接数={}",
                sessionId, sessionUserMap.size());
    }

    /**
     * 全量推送
     */
    public void sendAll(String message) {

        for (String sessionId : sessionUserMap.keySet()) {
            Session session = getSession(sessionId);
            send(session, message);
        }
    }

    /**
     * 按部门推送
     */
    public void sendToDept(Long deptId, String message) {

        Set<Session> sessions = deptSessionMap.get(deptId);
        if (sessions == null || sessions.isEmpty()) return;

        for (Session session : sessions) {
            send(session, message);
        }
    }

    /**
     * 完全异步发送
     */
    private void send(Session session, String message) {

        if (session == null) return;

        if (!session.isOpen()) {
            remove(session);
            return;
        }

        try {
            // 非阻塞发送
            session.getAsyncRemote().sendText(message);

        } catch (Exception e) {

            log.error("[WebSocket] 推送失败 sessionId={}", session.getId(), e);

            remove(session);

            try {
                session.close();
            } catch (IOException ex) {
                log.error("[WebSocket] 关闭失败", ex);
            }
        }
    }

    /**
     * 获取Session
     */
    private Session getSession(String sessionId) {

        for (Set<Session> sessions : deptSessionMap.values()) {
            for (Session s : sessions) {
                if (s.getId().equals(sessionId)) {
                    return s;
                }
            }
        }

        return null;
    }
}