package com.iqser.red.service.persistence.management.v1.processor.configuration;

import java.util.Collections;
import java.util.Optional;

import org.apache.tomcat.websocket.server.WsSci;
import org.springframework.boot.web.embedded.tomcat.TomcatContextCustomizer;
import org.springframework.boot.web.embedded.tomcat.TomcatServletWebServerFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.oauth2.server.resource.authentication.BearerTokenAuthenticationToken;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationToken;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

import com.knecon.fforesight.keycloakcommons.security.TenantAuthenticationManagerResolver;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;

@Slf4j
@Configuration
@EnableWebSocketMessageBroker
@RequiredArgsConstructor
public class WebSocketConfiguration implements WebSocketMessageBrokerConfigurer {

    private final TenantAuthenticationManagerResolver tenantAuthenticationManagerResolver;


    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {

        config.enableSimpleBroker("/topic");
        config.setApplicationDestinationPrefixes("/app");
    }


    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {

        registry.addEndpoint("/redaction-gateway-v1/websocket").setAllowedOrigins("*");
    }


    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {

        // https://docs.spring.io/spring-framework/reference/web/websocket/stomp/authentication-token-based.html
        registration.interceptors(new ChannelInterceptor() {
            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {

                StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
                if (StompCommand.CONNECT.equals(accessor.getCommand())) {
                    Optional.ofNullable(accessor.getNativeHeader("Authorization"))
                            .ifPresent(ah -> {
                                String bearerToken = ah.get(0).replace("Bearer ", "");
                                log.debug("Received bearer token {}", bearerToken);
                                AuthenticationManager authenticationManager = tenantAuthenticationManagerResolver.resolve(bearerToken);
                                JwtAuthenticationToken token = (JwtAuthenticationToken) authenticationManager.authenticate(new BearerTokenAuthenticationToken(bearerToken));
                                accessor.setUser(token);
                            });
                }
                return message;
            }
        });
    }


    @Bean
    public TomcatServletWebServerFactory tomcatContainerFactory() {

        TomcatServletWebServerFactory factory = new TomcatServletWebServerFactory();
        factory.setTomcatContextCustomizers(Collections.singletonList(tomcatContextCustomizer()));
        return factory;
    }


    @Bean
    public TomcatContextCustomizer tomcatContextCustomizer() {

        return context -> context.addServletContainerInitializer(new WsSci(), null);
    }

}
