[dependencies]
axum = {version = "0.7.9", features = ["ws"]}
tokio = {version = "1.41.1", features = ["full"]}

use axum::{
    response::Response,
    routing::get,
    Router,
    // Include the WebSocketUpgrade extractor:
    extract::ws::{WebSocket, WebSocketUpgrade},
};
use tokio;

// When a request comes in with the "Upgrade: websocket" header,
// you'll have `Some(ws)`. For regular HTTP requests, it will
// be `None`:
async fn handler(ws: Option<WebSocketUpgrade>) -> Response {
    if let Some(ws) = ws {
        // Handle WebSocket connection
        ws.on_upgrade(handle_socket)
    } else {
        // Handle HTTP request
        Response::new("Hello HTTP!".into())
    }
}

// Basic echo server example:
async fn handle_socket(mut socket: WebSocket) {
    while let Some(msg) = socket.recv().await {
        if let Ok(msg) = msg {
            if let Err(_) = socket.send(msg).await {
                break;
            }
        }
    }
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        // Pass the handler that includes the `WebSocketUpgrade` extractor to your route:
        .route("/", get(handler));

    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}