discord/src/chat.rs

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(Clone)]
pub struct ChatClient {
    tx: mpsc::UnboundedSender<ServerEvent>,
    is_idle: bool,
}

#[derive(Serialize, Clone)]
pub struct OnlineUser {
    pub user_id: Uuid,
    pub idle: bool,
}

#[derive(Default)]
pub struct ChatHub {
    // user_id -> client
    clients: RwLock<HashMap<Uuid, ChatClient>>,
}

#[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,
    },
    UserPresence {
        user_id: Uuid,
        online: bool,
        idle: bool,
    },
}

#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ClientEvent {
    // Currently no interactive client events for the general chat WS
    Ping,
    SetIdleStatus { is_idle: bool },
}

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, ChatClient { tx, is_idle: false });
    }

    pub async fn remove_client(&self, user_id: Uuid) {
        let mut clients = self.clients.write().await;
        clients.remove(&user_id);
    }

    pub async fn get_online_users(&self) -> Vec<OnlineUser> {
        let clients = self.clients.read().await;
        clients
            .iter()
            .map(|(id, client)| OnlineUser {
                user_id: *id,
                idle: client.is_idle,
            })
            .collect()
    }

    pub async fn broadcast_all(&self, event: ServerEvent) {
        let clients = self.clients.read().await;
        for client in clients.values() {
            let _ = client.tx.send(event.clone());
        }
    }

    pub async fn broadcast_to_user(&self, user_id: Uuid, event: ServerEvent) {
        let clients = self.clients.read().await;
        if let Some(client) = clients.get(&user_id) {
            let _ = client.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(client) = clients.get(&user_id) {
                let _ = client.tx.send(event.clone());
            }
        }
    }

    pub async fn set_idle_status(&self, user_id: Uuid, is_idle: bool) {
        {
            let mut clients = self.clients.write().await;
            if let Some(client) = clients.get_mut(&user_id) {
                client.is_idle = is_idle;
            }
        }
        self.broadcast_all(ServerEvent::UserPresence {
            user_id,
            online: true,
            idle: is_idle,
        })
        .await;
    }
}

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;
    state
        .chat
        .broadcast_all(ServerEvent::UserPresence {
            user_id,
            online: true,
            idle: false,
        })
        .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 {
        match msg {
            Message::Close(_) => break,
            Message::Text(text) => {
                if let Ok(ClientEvent::SetIdleStatus { is_idle }) =
                    serde_json::from_str::<ClientEvent>(&text)
                {
                    state.chat.set_idle_status(user_id, is_idle).await;
                }
            }
            _ => {}
        }
    }

    send_task.abort();
    state.chat.remove_client(user_id).await;
    state
        .chat
        .broadcast_all(ServerEvent::UserPresence {
            user_id,
            online: false,
            idle: false,
        })
        .await;
}