rustwave 0.1.0

use std::{future::Future, pin::Pin, sync::Arc};

use dashmap::DashMap;
use serde::Serialize;
use serde_json::Value;
use tokio::sync::mpsc::Sender;
use tokio_tungstenite::tungstenite::Message;
use uuid::Uuid;

use crate::error::Result;

use crate::peer::peer_state::PeerState;
use crate::{Payload, PeerContext, server::server_handler::ServerStates};

type NamespaceInnerType = Arc<
    DashMap<
        String,
        Arc<dyn Fn(Payload, PeerContext) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send + Sync>,
    >,
>;

/// Represents a single namespace within the server.
///
/// Each namespace maintains its own set of connected clients, rooms, and
/// event handlers, completely isolated from other namespaces. Clients
/// connected to one namespace cannot see or interact with clients,
/// events, or broadcasts belonging to another.
///
/// A [`Namespace`] is a lightweight, cloneable handle: cloning it does
/// not duplicate its state, it simply creates another reference to the
/// same underlying data (via `Arc`). This makes it safe and cheap to
/// share across tasks and threads.
///
/// # Example
///
/// ```no_run
/// use rustwave::{Server, Result};
///
/// #[tokio::main]
/// async fn main() -> Result<()> {
///     let server = Server::bind("127.0.0.1:8080");
///     let chat = server.namespace("/chat");
///
///     chat.on("message", |data, ctx| async move {
///         println!("Client {} say: {}", ctx.get_socket_id(), data);
///     });
///
///     server.run().await?;    
///     Ok(())
/// }
///
///
/// ```
pub struct Namespace {
    // Shared runtime state for this namespace: connected clients, room
    // membership, and the logic used for broadcasting/emitting messages.
    // Wrapped in `Arc` so all clones of this `Namespace` refer to the
    // same underlying state.
    state: Arc<ServerStates>,

    // Event handlers registered on this namespace.
    //
    // Key   — the event name (e.g. `"connection"`, `"message"`, `"disconnect"`)
    // Value — the handler invoked whenever that event is received
    inner: NamespaceInnerType,
}

impl Namespace {
    pub(crate) fn new() -> Self {
        Self {
            state: Arc::new(ServerStates::new()),
            inner: Arc::new(DashMap::new()),
        }
    }

    pub(crate) fn create_peer(&self, sender: Sender<Message>) -> PeerContext {
        let client = Arc::new(PeerState::new(sender, self.state.clone()));
        PeerContext::new(client)
    }

    pub(crate) fn handler(&self, event: &str, payload: Value, ctx: PeerContext) {
        if let Some(handler) = self.inner.get(event).map(|v| v.clone()) {
            let payload = Payload::new(payload);
            tokio::spawn(handler(payload, ctx));
        }
    }

    pub(crate) fn insert_client(&self, id: Uuid, sender: Sender<Message>) {
        self.state.insert_client(id, sender);
    }

    pub(crate) fn remove_client(&self, id: &Uuid) {
        self.state.remove_client(id);
    }

    pub(crate) fn remove_room(&self, room: &str, id: &Uuid) {
        self.state.remove_room(room, id);
    }

    /// Registers an event handler for this namespace.
    ///
    /// Registering a handler for an event that already has one replaces the
    /// previous handler.
    ///
    /// # Note
    ///
    /// The events `"connection"` and `"disconnect"` are fired automatically
    /// by the server when a client connects or disconnects, and can also be
    /// handled here like any other event.
    ///
    /// # Arguments
    ///
    /// * `event` - The name of the event to handle.
    /// * `handler` - The function to invoke when the event is received.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     // Register a handler for the event
    ///     chat.on("message", |data, ctx| async move {
    ///         println!("Client {} say: {}", ctx.get_socket_id(), data);
    ///     });
    /// # }
    /// ```
    pub fn on<F, Fut>(&self, event: &str, handler: F)
    where
        F: Fn(Payload, PeerContext) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        self.inner.insert(
            event.to_string(),
            Arc::new(move |data, ctx| Box::pin(handler(data, ctx))),
        );
    }

    /// Removes the event handler registered for `event` on this namespace.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    ///
    /// # #[tokio::main]
    /// # async fn main() {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     chat.on("message", |data, ctx| async move {
    ///         println!("Client {} say: {}", ctx.get_socket_id(), data);
    ///     });
    ///
    ///     // Remove the handler   
    ///     chat.off("message");
    /// # }
    /// ```
    pub fn off(&self, event: &str) {
        self.inner.remove(event);
    }

    /// Broadcasts an event to every client connected to this namespace.
    ///
    /// `data` is serialized to JSON and sent to all currently connected
    /// clients in this namespace. Clients connected to other namespaces will
    /// not receive this event.
    ///
    /// # Errors
    ///
    /// Returns an error if `data` fails to serialize.
    ///
    /// # Arguments
    ///
    /// * `event` - The name of the event to broadcast.
    /// * `data` - The data to broadcast.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::{Server, Result};
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     chat.emit("Event", "Payload!").await?;
    ///     Ok(())
    /// }
    /// ```
    pub async fn emit<T: Serialize>(&self, event: &str, data: T) -> Result<()> {
        self.state.broadcast(event, data).await?;
        Ok(())
    }

    /// Broadcasts an event to every client that has joined `room`, within
    /// this namespace.
    ///
    /// `data` is serialized to JSON and sent to all clients currently in
    /// `room`. Clients not in this room, or in a different namespace, will
    /// not receive this event. If the room doesn't exist or has no members,
    /// this is a no-op.
    ///
    /// # Errors
    ///
    /// Returns an error if `data` fails to serialize.
    ///
    /// # Arguments
    ///
    /// * `room` - The name of the room to broadcast to.
    /// * `event` - The name of the event to broadcast.
    /// * `data` - The data to broadcast.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    ///
    /// # #[tokio::main]
    /// # async fn example() -> rustwave::Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     chat.emit_to_room("room-1", "message", "Hello, room 1!").await?;
    ///     # Ok(())
    /// # }
    /// ```
    pub async fn emit_to_room<T: Serialize>(&self, room: &str, event: &str, data: T) -> Result<()> {
        self.state.room_broadcast(room, event, data).await?;
        Ok(())
    }

    /// Sends an event directly to a single client, identified by `socket_id`.
    ///
    /// `message` is serialized to JSON and sent only to the client matching
    /// `socket_id`. Unlike [`Namespace::emit`] and [`Namespace::emit_to_room`],
    /// this does not broadcast — only the targeted client receives it.
    ///
    /// # Errors
    ///
    /// Returns an error if `message` fails to serialize and ConnectionClosed, or if no client
    /// with the given `socket_id` is connected to this namespace.
    ///
    /// # Arguments
    ///
    /// * `socket_id` - The ID of the client to send the event to.
    /// * `event` - The name of the event to send.
    /// * `message` - The data to send.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    /// use uuid::Uuid;
    ///
    /// # #[tokio::main]
    /// # async fn example() -> rustwave::Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     let socket_id = Uuid::new_v4();
    ///
    ///     chat.emit_to(&socket_id, "private-message", "Hello there!").await?;
    ///     # Ok(())
    /// # }
    /// ```
    pub async fn emit_to<T: Serialize>(
        &self,
        socket_id: &Uuid,
        event: &str,
        message: T,
    ) -> Result<()> {
        self.state.emit_to(socket_id, event, message).await?;
        Ok(())
    }

    /// Checks whether a client with the given `socket_id` is currently
    /// connected to this namespace.
    ///
    /// Returns `true` if the client is connected, `false` otherwise (including
    /// if the client disconnected or `socket_id` never existed in this namespace).
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    /// use uuid::Uuid;
    ///
    /// # #[tokio::main]
    /// # async fn example() {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     let socket_id = Uuid::new_v4();
    ///
    ///     if chat.is_connected(&socket_id) {
    ///         println!("Client is still connected");
    ///     }
    /// # }
    /// ```
    pub fn is_connected(&self, socket_id: &Uuid) -> bool {
        self.state.is_connected(socket_id)
    }

    /// Returns the number of clients currently connected to this namespace.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    ///
    /// # #[tokio::main]
    /// async fn example() {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     println!("Connected clients: {}", chat.get_client_count());
    /// }
    /// ```
    pub fn get_client_count(&self) -> usize {
        self.state.get_client_count()
    }

    /// Forcibly disconnects a single client, identified by `socket_id`.
    ///
    /// Sends a close frame to the client and terminates its connection.
    /// This will trigger the client's `"disconnect"` event handler, if one
    /// is registered.
    ///
    /// # Errors
    ///
    /// Returns an error if no client with the given `socket_id` is connected
    /// to this namespace.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    /// use uuid::Uuid;
    ///
    /// # #[tokio::main]
    /// # async fn example() -> rustwave::Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     let socket_id = Uuid::new_v4();
    ///
    ///     chat.disconnect(&socket_id).await?;
    ///     # Ok(())
    /// # }
    /// ```
    pub async fn disconnect(&self, socket_id: &Uuid) -> Result<()> {
        self.state.disconnect(socket_id).await?;
        Ok(())
    }

    /// Forcibly disconnects every client currently connected to this namespace.
    ///
    /// Sends a close frame to each connected client and terminates their
    /// connections. This will trigger each client's `"disconnect"` event
    /// handler, if one is registered. Clients connected to other namespaces
    /// are not affected.
    ///
    /// # Errors
    ///
    /// Returns an error if the disconnect could not be dispatched.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use rustwave::Server;
    ///
    /// # #[tokio::main]
    /// # async fn example() -> rustwave::Result<()> {
    ///     let server = Server::bind("127.0.0.1:8080");
    ///     let chat = server.namespace("/chat");
    ///
    ///     chat.disconnect_all().await?;
    ///     # Ok(())
    /// # }
    /// ```
    pub async fn disconnect_all(&self) -> Result<()> {
        self.state.disconnect_all().await?;
        Ok(())
    }
}

impl Clone for Namespace {
    fn clone(&self) -> Self {
        Self {
            state: self.state.clone(),
            inner: self.inner.clone(),
        }
    }
}