poems 0.0.6

use futures_util::{SinkExt, StreamExt};
use poem::{
    // handler,
    web::websocket::{Message, WebSocket},
    web::FromRequest,
    // web::Path,
    Endpoint,
    IntoResponse,
    Request,
    Response,
    Result,
};

pub async fn ws_middleware<E>(next: E, req: Request) -> Result<Response>
where
    E: Endpoint,
{
    let path = req.uri().path();
    let path = path.to_string();
    println!("request: {}", path);
    let resp: Response = if path.ends_with("ws") {
        // let res = hello_ws.call(req).await.unwrap();
        let (req, mut body) = req.split();
        let ws = WebSocket::from_request(&req, &mut body).await.unwrap();
        let res = hello_ws(ws, path).await;
        println!("is ws");
        res.into_response()
    } else {
        let res = next.call(req).await.unwrap();
        println!("not ws");
        res.into_response()
    };
    Ok(resp)
}

// #[handler]
// pub fn hello_ws(Path(path): Path<String>, ws: WebSocket) -> impl IntoResponse {
pub async fn hello_ws(ws: WebSocket, path: String) -> impl IntoResponse {
    ws.on_upgrade(|socket| async move {
        let (mut tx, rx) = socket.split();
        let msg = Message::Text(format!("Hello, WebSocket {}", &path).into());
        // let msg = Message::Text("Hello, WebSocket".into());
        let _ = tx.send(msg).await;
        let _ = rx
            .for_each(|msg| async move {
                println!("Received message: {:?}", msg);
            })
            .await;
    })
}