package com.changhu.support.ws.interceptor;

import cn.hutool.core.lang.func.LambdaUtil;
import cn.hutool.core.util.StrUtil;
import com.changhu.common.utils.UserUtil;
import com.changhu.support.mybatisplus.pojo.entity.BaseEntity;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * author: luozhun
 * desc: webSocket握手拦截器
 * createTime: 2023/8/28 11:52
 */
@Slf4j
@Component
public class WsAuthenticationInterceptor implements HandshakeInterceptor {

    /**
     * ws如果在请求后拼接token的话太长了 这里用他提供的子协议进行传输
     */
    public static final String SEC_WEBSOCKET_PROTOCOL = "Sec-WebSocket-Protocol";

    public static final String attributeName = LambdaUtil.getFieldName(BaseEntity::getSnowFlakeId);

    /**
     * 初次握手访问前 看需不需要做权限认证
     *
     * @param request    request
     * @param response   response
     * @param wsHandler  wsHandler
     * @param attributes attributes
     * @return true or false
     */
    @Override
    public boolean beforeHandshake(@NotNull ServerHttpRequest request,
                                   @NotNull ServerHttpResponse response,
                                   @NotNull WebSocketHandler wsHandler,
                                   @NotNull Map<String, Object> attributes) {
        HttpServletRequest req = ((ServletServerHttpRequest) request).getServletRequest();
        //取出子协议的token
        String token = req.getHeader(SEC_WEBSOCKET_PROTOCOL);
        attributes.put(attributeName, UserUtil.getUserIdByToken(token));
        return true;
    }

    /**
     * 初次握手访问后，将前端自定义协议头Sec-WebSocket-Protocol原封不动返回回去，否则会报错
     *
     * @param request   the current request
     * @param response  the current response
     * @param wsHandler the target WebSocket handler
     * @param exception an exception raised during the handshake, or {@code null} if none
     */
    @Override
    public void afterHandshake(@NotNull ServerHttpRequest request,
                               @NotNull ServerHttpResponse response,
                               @NotNull WebSocketHandler wsHandler,
                               Exception exception) {
        HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
        HttpServletResponse httpResponse = ((ServletServerHttpResponse) response).getServletResponse();
        if (StrUtil.isNotEmpty(httpRequest.getHeader(SEC_WEBSOCKET_PROTOCOL))) {
            httpResponse.addHeader(SEC_WEBSOCKET_PROTOCOL, httpRequest.getHeader(SEC_WEBSOCKET_PROTOCOL));
        }
    }
}
