package com.njzscloud.common.ws.config;

import com.njzscloud.common.ws.support.ImpartialWebSocketHandler;
import com.njzscloud.common.ws.support.TokenHandshakeInterceptor;
import lombok.RequiredArgsConstructor;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.unit.DataSize;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.standard.ServletServerContainerFactoryBean;

import java.time.Duration;

@Configuration
@EnableWebSocket
@RequiredArgsConstructor
@EnableConfigurationProperties(WebsocketProperties.class)
public class WebSocketAutoConfiguration implements WebSocketConfigurer {

    private final WebsocketProperties websocketProperties;

    // websocket-heartbeat-js

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        String path = websocketProperties.getPath();
        registry.addHandler(new ImpartialWebSocketHandler(), path)
            .addInterceptors(new TokenHandshakeInterceptor())
            .setAllowedOrigins("*");
    }

    @Bean
    public ServletServerContainerFactoryBean createWebSocketContainer() {
        Duration asyncSendTimeout = websocketProperties.getAsyncSendTimeout();
        Duration maxSessionIdleTimeout = websocketProperties.getMaxSessionIdleTimeout();
        DataSize maxTextMessageBufferSize = websocketProperties.getMaxTextMessageBufferSize();
        DataSize maxBinaryMessageBufferSize = websocketProperties.getMaxBinaryMessageBufferSize();
        ServletServerContainerFactoryBean container = new ServletServerContainerFactoryBean();

        if (asyncSendTimeout != null) {
            long seconds = asyncSendTimeout.getSeconds();
            container.setAsyncSendTimeout(seconds * 1000);
        }

        if (maxSessionIdleTimeout != null) {
            long seconds = maxSessionIdleTimeout.getSeconds();
            container.setMaxSessionIdleTimeout(seconds * 1000);
        }

        if (maxTextMessageBufferSize != null) {
            long size = maxTextMessageBufferSize.toBytes();
            container.setMaxTextMessageBufferSize((int) size);
        }

        if (maxBinaryMessageBufferSize != null) {
            long size = maxBinaryMessageBufferSize.toBytes();
            container.setMaxBinaryMessageBufferSize((int) size);
        }

        return container;
    }
}
