@Component
public class WebSocketChannelInterceptor implements ChannelInterceptor {

    private final TokenUtils tokenUtils;

    private final AuthenticationManager websocketAuthenticationManager;

    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketChannelInterceptor.class);

    public WebSocketChannelInterceptor(TokenUtils tokenUtils, AuthenticationManager websocketAuthenticationManager) {
        this.tokenUtils = tokenUtils;
        this.websocketAuthenticationManager = websocketAuthenticationManager;
    }

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {

        LOGGER.info("WEBSOCKETCHANNELINTERCEPTOR -> "+message.toString());

        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);

        if (accessor != null && StompCommand.CONNECT.equals(accessor.getCommand())) {
            List<String> headers = accessor.getNativeHeader(AUTHORIZATION);
            accessor.setUser(websocketAuthenticationManager.authenticate(new JWTTokenAuthentication(tokenUtils.resolveToken(headers != null ? headers.get(0) : null))));
        }

        return message;
    }

    @Override
    public boolean preReceive(MessageChannel channel) {
        LOGGER.info("preReceive");
        return true;
    }

}

@Component("websocketAuthenticationManager")
@AllArgsConstructor
public class WebsocketAuthenticationManager implements AuthenticationManager {

    private final TokenUtils tokenUtils;

    private final ClientRepository clientRepository;

    private final OperatorRepository operatorRepository;

    @Override
    @Transactional(readOnly = true)
    public Authentication authenticate(Authentication authentication) throws AuthenticationException {

        JWTTokenAuthentication jwtTokenAuthentication = (JWTTokenAuthentication) authentication;

        String token = jwtTokenAuthentication.getToken();

        if(tokenUtils.validateToken(token)) {

            RoleType role = RoleType.get(tokenUtils.get("role", token));

            if(role == null) {
                throw throwBadCredentialsException("Unknown role");
            }

            switch (role) {
                case OPERATOR:

                    Optional<OperatorEntity> operator = operatorRepository.findByLogin(tokenUtils.get("login", token));

                    if(operator.isPresent()) {
                        return new WebsocketOperatorDetails(token, OperatorDetails.builder().role(role.toString()).operator(operator.get()).build());
                    } else {
                        throw throwBadCredentialsException("Operator not found");
                    }

                case CLIENT:

                    Optional<ClientEntity> client = clientRepository.findByHash(tokenUtils.get("hash", token));

                    if(client.isPresent()) {
                        return new WebsocketClientDetails(token, ClientDetails.builder().role(role.toString()).client(client.get()).build());
                    } else {
                        throw throwBadCredentialsException("Client not found");
                    }
                case OWNER:
                    throw throwBadCredentialsException("Access denied");
                default:
                    throw throwBadCredentialsException("Unknown role");
            }

        } else {
            throw throwBadCredentialsException("Invalid token");
        }

    }

    private MessagingException throwBadCredentialsException(String message) {

        Map<String, Object> headers = new HashMap<>();

        headers.put("status", StatusType.UNAUTHORIZED);
        headers.put("message", message);

        return new MessagingException(new MutableMessage<>(new MessageHeaders(headers)));
    }

}