package com.infoloop.tianting.intercepter;

import cn.dev33.satoken.stp.StpUtil;
import cn.hutool.core.convert.Convert;
import com.infoloop.tianting.constant.CommonConstants;
import com.infoloop.tianting.enums.LoginSourceEnum;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@SuppressWarnings("all")
public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) {
        final var authorization = request.getHeaders().getFirst(CommonConstants.AUTHORIZATION);
        if (StringUtils.isEmpty(authorization)) {
            return false;
        }
        final var loginSource = Convert.toInt(StpUtil.getExtra(CommonConstants.LOGIN_SOURCE));
        final var loginSourceEnum = LoginSourceEnum.getByValue(loginSource);
        attributes.put(CommonConstants.USER_TYPE, loginSourceEnum.getUserType().name());
        attributes.put(CommonConstants.USER_ID, StpUtil.getLoginIdAsString());
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
    }
}