# Stomp websocket setting

package com.alp.config;

import com.alp.jwt.exception.JwtTokenMalformedException;
import com.alp.jwt.util.JwtUtil;
import com.alp.util.CacheUtil;
import com.nimbusds.jwt.JWTClaimsSet;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.connection.MessageListener;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.listener.PatternTopic;
import org.springframework.data.redis.listener.RedisMessageListenerContainer;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.*;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import javax.servlet.http.HttpServletRequest;
import java.security.Principal;
import java.util.List;
import java.util.Map;

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig extends AbstractWebSocketMessageBrokerConfigurer/*AbstractSecurityWebSocketMessageBrokerConfigurer*/ implements WebSocketMessageBrokerConfigurer {
    private static final Logger log = LoggerFactory.getLogger(WebSocketConfig.class);

    @Autowired
    private JwtUtil jwtUtil;

    @Autowired
    private CacheUtil cacheUtil;

    @Override
    // 實現WebSocketMessageBrokerConfigurer中的此方法，配置消息代理(broker)
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker(
                "/DSN001", "/DSN002"); // 啟用SimpleBroker，使得訂閱到此"DSN001"、"DSN002"前綴的客戶端可以收到工作站通知.
    }

    @Override
    // 用來註冊Endpoint，"DSN001"、"DSN002"即為客戶端嘗試建立連接的地址。
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/DSN001", "/DSN002").setAllowedOrigins("*").setHandshakeHandler(handshakeHandler()).withSockJS();
    }


    /*@Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        super.configureClientInboundChannel(registration);
        registration.setInterceptors(new ChannelInterceptorAdapter() {
            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor =
                        MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

                log.info("accessor.getCommand():{}", accessor.getCommand());
                if (StompCommand.CONNECT.equals(accessor.getCommand())) {
                    //accessor.setUser(user);
                }
                return message;
            }
        });
    }*/
    
    @Override
    public void configureWebSocketTransport(WebSocketTransportRegistration registration) {
        super.configureWebSocketTransport(registration);
        registration.addDecoratorFactory(new WebSocketHandlerDecoratorFactory() {
            @Override
            public WebSocketHandler decorate(WebSocketHandler handler) {
                return new WebSocketHandlerDecorator(handler) {
                    @Override
                    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
                        super.afterConnectionEstablished(session);
                        log.info("afterConnectionEstablished:{}", session.getAttributes());
                        if (session.getAttributes().get("serviceTicket") != null) {
                            log.info("----force close----");
                            session.close(CloseStatus.POLICY_VIOLATION);
                        }
                    }
                };
            }
        });
    }

    @Bean
    public UserAuthHandshakeHandler handshakeHandler() {
        return new UserAuthHandshakeHandler();
    }

    public class UserAuthHandshakeHandler extends DefaultHandshakeHandler {
        @Override
        protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler, Map<String, Object> attributes) {
            // return request.getPrincipal(); // Default authentication

            // Handshake url parameter
            ServletServerHttpRequest servletServerHttpRequest = (ServletServerHttpRequest) request;
            HttpServletRequest httpRequest = servletServerHttpRequest.getServletRequest();
            //從連線的ws url的parameter取得參數
            String jwtToken = httpRequest.getParameter("token");

            if (StringUtils.isBlank(jwtToken)){
                throw new JwtTokenMalformedException( "no jwt token" );
            }
            JWTClaimsSet claims = jwtUtil.parseJwtToken( jwtToken );
            String serviceTicket = claims.getJWTID();
            attributes.put("serviceTicket",serviceTicket);
            log.info("request.getPrincipal(){}", request.getPrincipal());
            log.info("remoteAddress:{}", request.getRemoteAddress().getAddress().getHostAddress());
            log.info("get jwtToken:{}", jwtToken);
            return request.getPrincipal();
        }
    }

    @Override
    public void addArgumentResolvers(List<HandlerMethodArgumentResolver> argumentResolvers) {
        super.addArgumentResolvers(argumentResolvers);
    }

    @Bean(name = "dsn001PatternTopic")
    public PatternTopic dsn001PatternTopic() {
        return new PatternTopic("/DSN001/*");
    }

    @Bean(name = "dsn002PatternTopic")
    public PatternTopic dsn002PatternTopic() {
        return new PatternTopic("/DSN002/*");
    }

    /**
     * 設置一個處理 Redis Pub/Sub 消息的監聽器
     */
    @Bean
    public RedisMessageListenerContainer redisMessageListenerContainer(
            @Qualifier("wsRedisConnectionFactory") RedisConnectionFactory connectionFactory,
            @Qualifier("DsN001Mod") MessageListener dsN001Mod,
            @Qualifier("dsn001PatternTopic") PatternTopic dsn001PatternTopic,
            @Qualifier("DsN002Mod") MessageListener dsN002Mod,
            @Qualifier("dsn002PatternTopic") PatternTopic dsn002PatternTopic) {
        RedisMessageListenerContainer container = new RedisMessageListenerContainer();
        container.setConnectionFactory(connectionFactory);
        container.addMessageListener(dsN001Mod, dsn001PatternTopic);
        container.addMessageListener(dsN002Mod, dsn002PatternTopic);
        return container;
    }
}