g2jsonrt 0.1.0

use std::{collections::HashMap, net::SocketAddr};

use futures_util::{SinkExt, StreamExt};
use tokio::{
    net::{TcpListener, TcpStream},
    sync::mpsc::UnboundedSender,
};
use std::str::FromStr;

use tokio_tungstenite::{
    accept_async,
    tungstenite::{
        connect,
        handshake::server::{Request, Response},
    },
};

use axum::{
    body::HttpBody,
    extract::{
        ws::{Message, WebSocket, WebSocketUpgrade},
        State, TypedHeader,
    },
    headers,
    response::IntoResponse,
    routing::get,
    Router,
};

#[tokio::test]
async fn demo() {

    async fn handler(ws: WebSocketUpgrade) -> axum::response::Response {
        ws.on_upgrade(|socket| onhander(socket, |h, s, tx: tokio::sync::mpsc::UnboundedSender<Message>, hash| {
            tx.send(Message::Text("123".to_string()));
        }) )
    }
    let app = Router::new().route(
        "/ws",
        get(handler));
    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
    axum::Server::bind(&addr)
        .serve(app.into_make_service())
        .await
        .unwrap();

}

pub type Exec = fn(
    Header,
    &str,
    tokio::sync::mpsc::UnboundedSender<Message>,
    &mut std::collections::HashMap<i32, tokio::sync::mpsc::UnboundedSender<String>>,
);




pub async fn onhander(
    ws_stream: WebSocket,
    exec: fn(
        Header,
        &str,
        tokio::sync::mpsc::UnboundedSender<Message>,
        &mut std::collections::HashMap<i32, tokio::sync::mpsc::UnboundedSender<String>>,
    ),
) {
    match accept_connection(ws_stream, exec).await {
        Ok(()) => (),
        Err(e) => println!("{}", e),
    };
    ()
}

async fn accept_connection(
    ws_stream: WebSocket,
    exec: fn(
        Header,
        &str,
        tokio::sync::mpsc::UnboundedSender<Message>,
        &mut std::collections::HashMap<i32, tokio::sync::mpsc::UnboundedSender<String>>,
    ),
) -> Result<(), &'static str> {
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();

    let mut upstreams: HashMap<i32, UnboundedSender<String>> = HashMap::new();

    let (mut outgoing, mut incoming) = ws_stream.split();
    tokio::spawn(async move {
        while let Some(msg) = rx.recv().await {
            outgoing.send(msg).await;
        }
    });
    while let Some(msg) = incoming.next().await {
        let msg = match msg {
            Ok(msg) => msg,
            Err(e) => return Err("bad request for websocket"),
        };
        match msg {
            Message::Close(_) => return Err("connection closed"),
            Message::Binary(_) => return Err("bad request is binary"),
            Message::Ping(data) => {
                tx.send(Message::Pong(data)).map_err(|_| "send failed")?;
            }
            Message::Pong(_) => {
                continue;
            }
            Message::Text(text) => {
                let (header, body) = text.split_once("\n").ok_or_else(|| "bad request split \\n failed")?;
                let header: Header = serde_json::from_str(header).map_err(|_| "bad request parse header failed")?;
                exec(header, body, tx.clone(), &mut upstreams);
            }
        }
    }
    Ok(())
}

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct Header<'a> {
    pub path: &'a str,
    pub id: i32,
    pub status: Option<&'a str>,
}

#[derive(serde::Serialize)]
pub struct ReturnMessage<'a, T: serde::Serialize> {
    id: i32,
    status: Option<&'a str>,
    msg: Option<T>,
}

pub fn make_return_message<T: serde::Serialize>(
    id: i32,
    status: Option<&str>,
    msg: Option<T>,
) -> Message {
    let message = serde_json::to_string(&ReturnMessage { id, status, msg }).expect("bug");
    Message::Text(message)
}