use crate::api::AppState;
use crate::domain::auth::Jwt;
use axum::{
    extract::{Query, State, ws::WebSocketUpgrade},
    http::Extensions,
    response::IntoResponse,
};
use serde::Deserialize;
use tower_http::request_id::RequestId;

#[derive(Debug, Deserialize)]
pub struct WsParams {
    token: String,
}

pub async fn websocket_handler(
    ws: WebSocketUpgrade,
    Query(params): Query<WsParams>,
    extensions: Extensions,
    State(state): State<AppState>,
) -> impl IntoResponse {
    let request_id = extensions
        .get::<RequestId>()
        .map_or_else(|| "unknown".to_string(), |id| id.header_value().to_str().unwrap_or_default().to_string());

    let jwt = Jwt::new(params.token);
    match state.auth_service.verify_token(&jwt) {
        Ok(user_id) => ws.on_upgrade(move |socket| {
            let service = state.gateway_service.clone();
            let shutdown = state.shutdown_rx.clone();
            async move {
                service.handle_socket(socket, user_id, request_id, shutdown).await;
            }
        }),
        Err(e) => {
            tracing::warn!(error = %e, "WebSocket handshake failed: invalid token");
            axum::http::StatusCode::UNAUTHORIZED.into_response()
        }
    }
}
