discord/src/voice.rs

use std::collections::HashMap;

use axum::extract::ws::{Message, WebSocket};
use futures_util::{SinkExt, StreamExt};
use serde::{Deserialize, Serialize};
use tokio::sync::{RwLock, mpsc};
use uuid::Uuid;

use crate::{AppState, db};

#[derive(Default)]
pub struct VoiceHub {
    rooms: RwLock<HashMap<Uuid, HashMap<Uuid, ClientHandle>>>,
}

#[derive(Clone)]
struct ClientHandle {
    display_name: String,
    is_sharing_video: bool,
    tx: mpsc::UnboundedSender<ServerEvent>,
}

#[derive(Serialize, Clone)]
pub struct VoiceParticipant {
    pub user_id: Uuid,
    pub display_name: String,
    pub is_sharing_video: bool,
}

#[derive(Serialize, Clone)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ServerEvent {
    Peers {
        peers: Vec<VoiceParticipant>,
    },
    PeerJoined {
        user_id: Uuid,
        display_name: String,
    },
    PeerLeft {
        user_id: Uuid,
    },
    VideoStatusChanged {
        user_id: Uuid,
        is_sharing_video: bool,
    },
    Signal {
        from_user_id: Uuid,
        kind: String,
        data: serde_json::Value,
    },
    Error {
        message: String,
    },
}

#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ClientEvent {
    Signal {
        to_user_id: Uuid,
        kind: String,
        data: serde_json::Value,
    },
    SetVideoStatus {
        is_sharing_video: bool,
    },
}

impl VoiceHub {
    pub async fn participants(&self, room_id: Uuid) -> Vec<VoiceParticipant> {
        let rooms = self.rooms.read().await;
        let Some(room) = rooms.get(&room_id) else {
            return Vec::new();
        };

        room.iter()
            .map(|(user_id, handle)| VoiceParticipant {
                user_id: *user_id,
                display_name: handle.display_name.clone(),
                is_sharing_video: handle.is_sharing_video,
            })
            .collect()
    }

    async fn join(
        &self,
        room_id: Uuid,
        user_id: Uuid,
        display_name: String,
        tx: mpsc::UnboundedSender<ServerEvent>,
    ) -> Vec<VoiceParticipant> {
        let mut rooms = self.rooms.write().await;
        let room = rooms.entry(room_id).or_default();

        let peers = room
            .iter()
            .map(|(peer_id, peer)| VoiceParticipant {
                user_id: *peer_id,
                display_name: peer.display_name.clone(),
                is_sharing_video: peer.is_sharing_video,
            })
            .collect::<Vec<_>>();

        room.insert(
            user_id,
            ClientHandle {
                display_name: display_name.clone(),
                is_sharing_video: false,
                tx,
            },
        );

        for (peer_id, peer) in room.iter() {
            if *peer_id != user_id {
                let _ = peer.tx.send(ServerEvent::PeerJoined {
                    user_id,
                    display_name: display_name.clone(),
                });
            }
        }

        peers
    }

    async fn leave(&self, room_id: Uuid, user_id: Uuid) {
        let mut rooms = self.rooms.write().await;
        let Some(room) = rooms.get_mut(&room_id) else {
            return;
        };

        room.remove(&user_id);

        for peer in room.values() {
            let _ = peer.tx.send(ServerEvent::PeerLeft { user_id });
        }

        if room.is_empty() {
            rooms.remove(&room_id);
        }
    }

    pub async fn relay_signal(
        &self,
        room_id: Uuid,
        from_user_id: Uuid,
        to_user_id: Uuid,
        kind: String,
        data: serde_json::Value,
    ) {
        let rooms = self.rooms.read().await;
        let Some(room) = rooms.get(&room_id) else {
            return;
        };

        if let Some(target) = room.get(&to_user_id) {
            let _ = target.tx.send(ServerEvent::Signal {
                from_user_id,
                kind,
                data,
            });
        }
    }

    pub async fn set_video_status(&self, room_id: Uuid, user_id: Uuid, is_sharing_video: bool) {
        let mut rooms = self.rooms.write().await;
        let Some(room) = rooms.get_mut(&room_id) else {
            return;
        };

        if let Some(handle) = room.get_mut(&user_id) {
            handle.is_sharing_video = is_sharing_video;

            for (peer_id, peer) in room.iter() {
                if *peer_id != user_id {
                    let _ = peer.tx.send(ServerEvent::VideoStatusChanged {
                        user_id,
                        is_sharing_video,
                    });
                }
            }
        }
    }
}

pub async fn handle_socket(state: AppState, socket: WebSocket, room_id: Uuid, user_id: Uuid) {
    let Some(user) = db::get_user_by_id(&state.db, user_id).await.ok().flatten() else {
        return;
    };

    let (mut ws_sender, mut ws_receiver) = socket.split();
    let (tx, mut rx) = mpsc::unbounded_channel::<ServerEvent>();

    let peers = state
        .voice
        .join(room_id, user_id, user.display_name.clone(), tx.clone())
        .await;
    let _ = tx.send(ServerEvent::Peers { peers });

    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::Text(text) => {
                let parsed = serde_json::from_str::<ClientEvent>(&text);
                match parsed {
                    Ok(ClientEvent::Signal {
                        to_user_id,
                        kind,
                        data,
                    }) => {
                        state
                            .voice
                            .relay_signal(room_id, user_id, to_user_id, kind, data)
                            .await;
                    }
                    Ok(ClientEvent::SetVideoStatus { is_sharing_video }) => {
                        state
                            .voice
                            .set_video_status(room_id, user_id, is_sharing_video)
                            .await;
                    }
                    Err(err) => {
                        let _ = tx.send(ServerEvent::Error {
                            message: format!("invalid voice message: {err}"),
                        });
                    }
                }
            }
            Message::Close(_) => break,
            _ => {}
        }
    }

    send_task.abort();
    state.voice.leave(room_id, user_id).await;
}