argentor-gateway 1.4.7

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{mpsc, RwLock};
use uuid::Uuid;

/// Represents a connected WebSocket client.
#[derive(Debug)]
pub struct Connection {
    /// Unique connection identifier.
    pub id: Uuid,
    /// Session this connection belongs to.
    pub session_id: Uuid,
    /// Channel sender for outbound messages.
    pub tx: mpsc::UnboundedSender<String>,
}

/// Manages active WebSocket connections.
pub struct ConnectionManager {
    connections: RwLock<HashMap<Uuid, Connection>>,
}

impl ConnectionManager {
    /// Create a new connection manager wrapped in an `Arc`.
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            connections: RwLock::new(HashMap::new()),
        })
    }

    /// Register a new connection.
    pub async fn add(&self, conn: Connection) {
        let id = conn.id;
        self.connections.write().await.insert(id, conn);
        tracing::info!(connection_id = %id, "Connection added");
    }

    /// Remove a connection by ID.
    pub async fn remove(&self, id: Uuid) {
        self.connections.write().await.remove(&id);
        tracing::info!(connection_id = %id, "Connection removed");
    }

    /// Send a message to all connections belonging to the given session.
    pub async fn send_to_session(&self, session_id: Uuid, message: &str) {
        let conns = self.connections.read().await;
        for conn in conns.values() {
            if conn.session_id == session_id {
                let _ = conn.tx.send(message.to_string());
            }
        }
    }

    /// Return the number of active connections.
    pub async fn connection_count(&self) -> usize {
        self.connections.read().await.len()
    }

    /// Return the set of unique session IDs across all active connections.
    pub async fn session_ids(&self) -> Vec<Uuid> {
        let conns = self.connections.read().await;
        let mut ids: Vec<Uuid> = conns.values().map(|c| c.session_id).collect();
        ids.sort();
        ids.dedup();
        ids
    }

    /// Broadcast a message to all connected clients.
    pub async fn broadcast(&self, message: &str) {
        let conns = self.connections.read().await;
        for conn in conns.values() {
            let _ = conn.tx.send(message.to_string());
        }
    }
}

impl Default for ConnectionManager {
    fn default() -> Self {
        Self {
            connections: RwLock::new(HashMap::new()),
        }
    }
}