package org.xyzh.workcase.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.xyzh.common.auth.contants.AuthContants;

import java.util.Map;

/**
 * WebSocket握手拦截器
 * 从网关传递的请求头中获取已验证的用户信息
 */
@Component
public class WebSocketAuthInterceptor implements HandshakeInterceptor {

    private static final Logger log = LoggerFactory.getLogger(WebSocketAuthInterceptor.class);

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {

        log.info("WebSocket握手开始，URI: {}", request.getURI());

        try {
            // 从网关传递的请求头中获取用户ID（网关已完成token验证）
            String userId = extractUserIdFromRequest(request);

            if (!StringUtils.hasText(userId)) {
                log.warn("WebSocket握手失败：请求头中未找到用户ID，网关认证可能失败");
                return false;
            }

            // 可选：也获取token，供后续使用
            String token = extractTokenFromRequest(request);

            // 将用户信息存储到WebSocket会话属性中，供后续使用
            attributes.put("userId", userId);
            if (StringUtils.hasText(token)) {
                attributes.put("token", token);
            }

            log.info("WebSocket握手成功，userId: {}", userId);
            return true;

        } catch (Exception e) {
            log.error("WebSocket握手异常", e);
            return false;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                              WebSocketHandler wsHandler, Exception exception) {
        if (exception != null) {
            log.error("WebSocket握手后发生异常", exception);
        }
    }

    /**
     * 从网关传递的请求头中提取用户ID
     */
    private String extractUserIdFromRequest(ServerHttpRequest request) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;

            // 从网关传递的请求头获取（标准微服务架构方式）
            String userId = servletRequest.getHeaders().getFirst(AuthContants.USER_ID_ATTRIBUTE);
            if (StringUtils.hasText(userId)) {
                log.debug("从网关请求头获取用户ID: {}", userId);
                return userId;
            }
        }

        log.warn("未能从请求头中获取用户ID");
        return null;
    }

    /**
     * 从请求头中提取token（可选）
     */
    private String extractTokenFromRequest(ServerHttpRequest request) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;

            // 从网关传递的请求头获取
            String token = servletRequest.getHeaders().getFirst(AuthContants.TOKEN_ATTRIBUTE);
            if (StringUtils.hasText(token)) {
                return token;
            }

            // 也支持从Authorization头获取
            String authHeader = servletRequest.getHeaders().getFirst("Authorization");
            if (StringUtils.hasText(authHeader)) {
                if (authHeader.startsWith("Bearer ")) {
                    return authHeader.substring(7);
                }
                return authHeader;
            }
        }

        return null;
    }
}
