package com.ttstd.signaling.security;

import com.ttstd.signaling.model.AuthPrincipal;
import com.ttstd.signaling.service.AccountService;
import com.ttstd.signaling.service.DeviceIdentityService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
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.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.net.URI;
import java.util.List;
import java.util.Map;

/**
 * WebSocket 握手鉴权拦截器。
 *
 * <p>连接建立前完成身份认定，未通过者直接以 401 拒绝握手，杜绝匿名连接。
 * 认证通过后将 {@link AuthPrincipal} 写入会话属性，后续信令处理一律以此为准，
 * 客户端自报的 fromDeviceId / deviceType 不再被信任。
 *
 * <p>令牌传递优先级：
 * <ol>
 *   <li>{@code Sec-WebSocket-Protocol: signal.v1, auth.<token>}（推荐，不进访问日志）</li>
 *   <li>{@code Authorization: Bearer <token>}（非浏览器客户端可用）</li>
 *   <li>URL 查询参数 {@code ?token=}（兼容用途，会告警）</li>
 * </ol>
 */
@Component
public class AuthHandshakeInterceptor implements HandshakeInterceptor {

    private static final Logger logger = LoggerFactory.getLogger(AuthHandshakeInterceptor.class);

    /** 会话属性键：已鉴权主体 */
    public static final String ATTR_PRINCIPAL = "authPrincipal";
    /** 会话属性键：协商选择的子协议 */
    public static final String ATTR_SELECTED_PROTOCOL = "selectedProtocol";

    private static final String SUBPROTOCOL = "signal.v1";
    private static final String AUTH_PREFIX = "auth.";

    /** 单 IP 握手频率限制：60 秒内最多 30 次 */
    private static final int HANDSHAKE_LIMIT = 30;
    private static final long HANDSHAKE_WINDOW_SECONDS = 60;

    private final AccountService accountService;
    private final DeviceIdentityService deviceIdentityService;
    private final RateLimiter rateLimiter;

    public AuthHandshakeInterceptor(AccountService accountService,
                                    DeviceIdentityService deviceIdentityService,
                                    RateLimiter rateLimiter) {
        this.accountService = accountService;
        this.deviceIdentityService = deviceIdentityService;
        this.rateLimiter = rateLimiter;
    }

    @Override
    public boolean beforeHandshake(ServerHttpRequest request,
                                   ServerHttpResponse response,
                                   WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) {
        String clientIp = resolveClientIp(request);

        if (!rateLimiter.tryAcquire("ws:" + clientIp, HANDSHAKE_LIMIT, HANDSHAKE_WINDOW_SECONDS)) {
            logger.warn("握手请求过于频繁，已拒绝：ip={}", clientIp);
            response.setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
            return false;
        }

        String token = extractToken(request);
        if (token == null || token.isBlank()) {
            logger.warn("握手被拒绝：未携带令牌，ip={}", clientIp);
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;
        }

        AuthPrincipal principal = resolvePrincipal(token, clientIp);
        if (principal == null) {
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;
        }

        attributes.put(ATTR_PRINCIPAL, principal);
        // 若客户端使用子协议方式传递令牌，需回显子协议名完成协商
        if (requestsSubprotocol(request)) {
            attributes.put(ATTR_SELECTED_PROTOCOL, SUBPROTOCOL);
        }

        logger.info("握手鉴权通过：{} {} (deviceId={}, ip={})",
                principal.principalType(), principal.displayName(), principal.deviceId(), clientIp);
        return true;
    }

    /**
     * 依次尝试按用户令牌、设备令牌解析身份。
     * 两者均失败时统一返回 null，不向客户端区分失败原因。
     */
    private AuthPrincipal resolvePrincipal(String token, String clientIp) {
        try {
            return accountService.authenticate(token);
        } catch (AuthException userEx) {
            try {
                return deviceIdentityService.authenticate(token);
            } catch (AuthException deviceEx) {
                logger.warn("握手鉴权失败，ip={}，用户令牌校验：{}；设备令牌校验：{}",
                        clientIp, userEx.getMessage(), deviceEx.getMessage());
                return null;
            }
        }
    }

    private boolean requestsSubprotocol(ServerHttpRequest request) {
        List<String> protocols = request.getHeaders().get("Sec-WebSocket-Protocol");
        if (protocols == null) {
            return false;
        }
        for (String raw : protocols) {
            for (String part : raw.split(",")) {
                if (SUBPROTOCOL.equals(part.trim())) {
                    return true;
                }
            }
        }
        return false;
    }

    /**
     * 提取令牌，优先使用不会被写入访问日志的传递方式。
     */
    private String extractToken(ServerHttpRequest request) {
        // 1) Sec-WebSocket-Protocol: signal.v1, auth.<token>
        List<String> protocols = request.getHeaders().get("Sec-WebSocket-Protocol");
        if (protocols != null) {
            for (String raw : protocols) {
                for (String part : raw.split(",")) {
                    String value = part.trim();
                    if (value.startsWith(AUTH_PREFIX) && value.length() > AUTH_PREFIX.length()) {
                        return value.substring(AUTH_PREFIX.length());
                    }
                }
            }
        }

        // 2) Authorization: Bearer <token>
        String authorization = request.getHeaders().getFirst("Authorization");
        if (authorization != null && authorization.regionMatches(true, 0, "Bearer ", 0, 7)) {
            return authorization.substring(7).trim();
        }

        // 3) URL 查询参数（兼容旧客户端，存在写入日志的风险）
        URI uri = request.getURI();
        String query = uri.getQuery();
        if (query != null) {
            for (String pair : query.split("&")) {
                int idx = pair.indexOf('=');
                if (idx > 0 && "token".equals(pair.substring(0, idx))) {
                    logger.warn("客户端通过 URL 查询参数传递令牌，存在日志泄露风险，建议改用 Sec-WebSocket-Protocol");
                    return java.net.URLDecoder.decode(
                            pair.substring(idx + 1), java.nio.charset.StandardCharsets.UTF_8);
                }
            }
        }
        return null;
    }

    private String resolveClientIp(ServerHttpRequest request) {
        String forwarded = request.getHeaders().getFirst("X-Forwarded-For");
        if (forwarded != null && !forwarded.isBlank()) {
            int comma = forwarded.indexOf(',');
            return (comma > 0 ? forwarded.substring(0, comma) : forwarded).trim();
        }
        if (request instanceof ServletServerHttpRequest servletRequest) {
            return servletRequest.getServletRequest().getRemoteAddr();
        }
        return request.getRemoteAddress() == null
                ? "unknown"
                : request.getRemoteAddress().getAddress().getHostAddress();
    }

    @Override
    public void afterHandshake(ServerHttpRequest request,
                               ServerHttpResponse response,
                               WebSocketHandler wsHandler,
                               Exception exception) {
        // 无需处理
    }
}
