ui stuff

src/chat.rs

@@ -0,0 +1,88 @@
+use crate::AppState;
+use axum::extract::ws::{Message, WebSocket};
+use futures_util::{SinkExt, StreamExt};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use tokio::sync::{RwLock, mpsc};
+use uuid::Uuid;
+
+#[derive(Default)]
+pub struct ChatHub {
+    // user_id -> sender
+    clients: RwLock<HashMap<Uuid, mpsc::UnboundedSender<ServerEvent>>>,
+}
+
+#[derive(Serialize, Clone)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum ServerEvent {
+    MessageCreated {
+        channel_id: Uuid,
+        message: serde_json::Value,
+    },
+    DmCreated {
+        other_user_id: Uuid,
+        message: serde_json::Value,
+    },
+}
+
+#[derive(Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+enum ClientEvent {
+    // Currently no interactive client events for the general chat WS
+    Ping,
+}
+
+impl ChatHub {
+    pub async fn add_client(&self, user_id: Uuid, tx: mpsc::UnboundedSender<ServerEvent>) {
+        let mut clients = self.clients.write().await;
+        clients.insert(user_id, tx);
+    }
+
+    pub async fn remove_client(&self, user_id: Uuid) {
+        let mut clients = self.clients.write().await;
+        clients.remove(&user_id);
+    }
+
+    pub async fn broadcast_to_user(&self, user_id: Uuid, event: ServerEvent) {
+        let clients = self.clients.read().await;
+        if let Some(tx) = clients.get(&user_id) {
+            let _ = tx.send(event);
+        }
+    }
+
+    pub async fn broadcast_to_many(&self, user_ids: Vec<Uuid>, event: ServerEvent) {
+        let clients = self.clients.read().await;
+        for user_id in user_ids {
+            if let Some(tx) = clients.get(&user_id) {
+                let _ = tx.send(event.clone());
+            }
+        }
+    }
+}
+
+pub async fn handle_socket(state: AppState, socket: WebSocket, user_id: Uuid) {
+    let (mut ws_sender, mut ws_receiver) = socket.split();
+    let (tx, mut rx) = mpsc::unbounded_channel::<ServerEvent>();
+
+    state.chat.add_client(user_id, tx).await;
+
+    let send_task = tokio::spawn(async move {
+        while let Some(event) = rx.recv().await {
+            let Ok(payload) = serde_json::to_string(&event) else {
+                continue;
+            };
+            if ws_sender.send(Message::Text(payload.into())).await.is_err() {
+                break;
+            }
+        }
+    });
+
+    while let Some(Ok(msg)) = ws_receiver.next().await {
+        if let Message::Close(_) = msg {
+            break;
+        }
+    }
+
+    send_task.abort();
+    state.chat.remove_client(user_id).await;
+}