package com.njzscloud.common.mvc.ws.config;

import com.njzscloud.common.mvc.ws.support.TokenHandshakeInterceptor;
import com.njzscloud.common.mvc.ws.support.WebSocketChannelInterceptor;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

/**
 * Websocket 配置，使用 STOMP 协议
 */
@Configuration(proxyBeanMethods = false)
@RequiredArgsConstructor
@EnableWebSocketMessageBroker
@EnableConfigurationProperties(WebsocketProperties.class)
@ConditionalOnProperty(value = "websocket.enable", havingValue = "true")
public class WebsocketAutoConfiguration implements WebSocketMessageBrokerConfigurer, DisposableBean {

    private final WebsocketProperties websocketProperties;
    private final WebSocketChannelInterceptor webSocketChannelInterceptor = new WebSocketChannelInterceptor();
    private final TokenHandshakeInterceptor tokenHandshakeInterceptor = new TokenHandshakeInterceptor();

    /**
     * 用于心跳消息的线程池
     */
    private final ThreadPoolTaskScheduler tts = new ThreadPoolTaskScheduler();

    {

        tts.setPoolSize(1);
        tts.setThreadNamePrefix("ws-heartbeat-thread-");
        tts.initialize();
    }

    /**
     * 设置 Websocket
     */
    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint(websocketProperties.getEndpoint())
                .addInterceptors(tokenHandshakeInterceptor);
    }

    /**
     * 配置 STOMP 消息服务器
     */
    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        WebsocketProperties.StompProperties stomp = websocketProperties.getStomp();
        config.setApplicationDestinationPrefixes(stomp.getApplicationPrefixes())
                .enableSimpleBroker(stomp.getBroadcastPrefixes())
                .setTaskScheduler(tts);
    }

    /**
     * 消息拦截器
     */
    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(webSocketChannelInterceptor);
    }

    @Override
    public void destroy() throws Exception {
        tts.shutdown();
    }
}
