use futures::stream::SplitStream; use tokio::sync::Mutex; use futures::SinkExt; use tokio::net::{TcpListener, TcpStream}; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::{accept_async, tungstenite}; use futures::StreamExt; use std::sync::Arc; use serde_json::{Value, json}; use crate::messages::send_dm; use crate::rs_println; use crate::Args; type Sender = Arc, tungstenite::Message>>>>; type Receiver = Arc>>>>; static mut GLOBAL_SENDER: Option = None; static mut GLOBAL_RECEIVER: Option = None; static mut REPLY_HELLO: bool = false; pub static mut HAS_CONNECTED: bool = false; async fn set_sender(sender: Sender) { unsafe { GLOBAL_SENDER = Some(sender); } } async fn set_receiver(receiver: Receiver) { unsafe { GLOBAL_RECEIVER = Some(receiver); } } #[allow(static_mut_refs)] pub async fn send_msg(msg: &str) { unsafe { if let Some(sender) = &GLOBAL_SENDER { let mut sender = sender.lock().await; if let Some(s) = sender.as_mut() { s.send(tungstenite::Message::Text(msg.to_string().into())).await.unwrap(); } } } } #[allow(static_mut_refs)] pub async fn send_cmd_json(func_name: &str, func_args: Option) -> Option { unsafe { let Some(sender) = &GLOBAL_SENDER else { return None }; let mut sender = sender.lock().await; let Some(s) = sender.as_mut() else { return None }; let unw_args = if func_args.is_some() { func_args.unwrap() } else { json!([]) }; let json_str = format!( "json:{{\"type\": \"function\", \"value\":\"{}\", \"args\": {}}}", func_name, unw_args ); if s.send(tungstenite::Message::Text(json_str.into())).await.is_err() { return None; } let r = receive_response().await; rs_println!("Received from Python: [RESPONSE] {:?}", r); if r.is_none() { rs_println!("--- Response from Python is None!"); } return r; } } #[allow(static_mut_refs)] async fn receive_response() -> Option { unsafe { let Some(receiver) = &GLOBAL_RECEIVER else { return None }; let mut receiver = receiver.lock().await; let Some(r) = receiver.as_mut() else { return None }; let Some(Ok(msg)) = r.next().await else { return None }; let tungstenite::Message::Text(response) = msg else { return None }; if response.starts_with("json:") { return serde_json::from_str(&response[5..]).ok(); } else { return serde_json::from_str(&response).ok(); } } } pub async fn start(args: Args) { rs_println!("Starting local websocket..."); let ip = format!("127.0.0.1:{}", args.port); let listener = TcpListener::bind(&ip).await.unwrap(); rs_println!("WebSocket server running on ws://{}", ip); tokio::spawn(handle_connections(listener, args)); } async fn handle_connections(listener: TcpListener, args: Args) { while let Ok((stream, _)) = listener.accept().await { let ws_stream = accept_async(stream).await.unwrap(); let (sender, receiver) = ws_stream.split(); let sender_arc = Arc::new(Mutex::new(Some(sender))); let receiver_arc = Arc::new(Mutex::new(Some(receiver))); set_sender(sender_arc.clone()).await; set_receiver(receiver_arc.clone()).await; while let Some(Ok(msg)) = receiver_arc.lock().await.as_mut().unwrap().next().await { handle_message(msg, args.clone()).await; } } } async fn handle_message(msg: tungstenite::protocol::Message, args: Args) { match msg { tungstenite::Message::Text(text) => { rs_println!("Received from Python: {}", text); if text.starts_with("json:") { let t_json: Value = serde_json::from_str(&text[5..]).unwrap(); if t_json.get("error").is_some() { send_dm("Internal Python error!".to_string(), args).await; } } unsafe { if !REPLY_HELLO { send_msg("[Connection test] Hello from Rust!").await; REPLY_HELLO = true; HAS_CONNECTED = true; } } } tungstenite::Message::Binary(bytes) => { if args.dev && !args.noping { rs_println!("[Binary] from Python: {:?}", bytes); } } _ => { if args.dev && !args.noping { rs_println!("Received from Python: [UNKNOWN / OTHER]"); } } } }