package com.centreproject.common.websocket;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.Map;
import java.util.concurrent.*;

/**
 * @author zhonglei.xu
 */
public class MyWebSocketHandler extends TextWebSocketHandler {
    private static final Logger LOG = LoggerFactory.getLogger(MyWebSocketHandler.class);
    private static final Map<String, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();
    private final Map<String, BlockingQueue<String>> sessionQueues = new ConcurrentHashMap<>();
    private final RedisTemplate<String, String> redisTemplate;
    private final ExecutorService executor = new ThreadPoolExecutor(
            4,
            32,
            30, TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(1000),
            new ThreadPoolExecutor.CallerRunsPolicy()
    );

    public MyWebSocketHandler(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        SESSIONS.put(session.getId(), session);
        LOG.info("客户端连接: ID={}, IP={}",
                session.getId(),
                session.getRemoteAddress().getAddress().getHostAddress());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        SESSIONS.remove(session.getId());
        LOG.info("客户端断开: ID={}, 原因={}", session.getId(), status.getReason());
    }

    /**
     * 添加会话队列处理方法
     * @param jsonData jsonData
     */
    public void localBroadcast(String jsonData) {
        SESSIONS.forEach((sessionId, session) -> {
            sessionQueues
                    .computeIfAbsent(sessionId, k -> new LinkedBlockingQueue<>())
                    .offer(jsonData);
            processQueue(session);
        });
    }

    private void processQueue(WebSocketSession session) {
        executor.submit(() -> {
            try {
                BlockingQueue<String> queue = sessionQueues.get(session.getId());
                while (session.isOpen() && !queue.isEmpty()) {
                    String message = queue.poll();
                    synchronized (session) {  // 会话级同步
                        if (session.isOpen()) {
                            session.sendMessage(new TextMessage(message));
                        }
                    }
                }
            } catch (Exception e) {
                SESSIONS.remove(session.getId());
                sessionQueues.remove(session.getId());
            }
        });
    }


    /**
     全局广播（通过Redis发布，所有实例消费）
     * @param jsonData jsonData
     */
    public void broadcast(String jsonData) {
        redisTemplate.convertAndSend("websocket-broadcast", jsonData);
    }
}