gitdataai/lib/socketio/namespace.rs

use std::{
    collections::HashMap, collections::HashSet, future::Future, sync::Arc,
};

use serde::Serialize;
use serde_json::Value;
use tokio::sync::RwLock;

use crate::{
    adapter::{Adapter, BroadcastOptions},
    config::SocketIoConfig,
    error::Result,
    packet::EventPayload,
    server::{Inner, Namespace, NamespaceState, SocketIo, SocketIoBuilder},
    socket::{DisconnectReason, Socket},
};

impl SocketIoBuilder {
    pub fn config(mut self, config: SocketIoConfig) -> Self {
        self.config = config;
        self
    }

    pub fn adapter(mut self, adapter: Arc<dyn Adapter>) -> Self {
        self.adapter = adapter;
        self
    }

    pub fn build(self) -> SocketIo {
        let mut namespaces = HashMap::new();
        namespaces.insert("/".to_owned(), Arc::new(NamespaceState::default()));
        SocketIo {
            inner: Arc::new(Inner {
                config: self.config,
                sessions: RwLock::new(HashMap::new()),
                namespaces: RwLock::new(namespaces),
                adapter: self.adapter,
                next_ack_id: std::sync::atomic::AtomicU64::new(1),
            }),
        }
    }
}

impl Namespace {
    pub async fn on_connect<F, Fut>(&self, handler: F)
    where
        F: Fn(Socket) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let state = self.io.ensure_namespace(&self.name).await;
        *state.connect_handler.write().await =
            Some(Arc::new(move |socket| Box::pin(handler(socket))));
    }

    pub async fn on_disconnect<F, Fut>(&self, handler: F)
    where
        F: Fn(Socket, DisconnectReason) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let state = self.io.ensure_namespace(&self.name).await;
        *state.disconnect_handler.write().await =
            Some(Arc::new(move |socket, reason| {
                Box::pin(handler(socket, reason))
            }));
    }

    pub async fn use_middleware<F, Fut>(&self, middleware: F)
    where
        F: Fn(Socket, Option<Value>) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<()>> + Send + 'static,
    {
        let state = self.io.ensure_namespace(&self.name).await;
        state
            .middleware
            .write()
            .await
            .push(Arc::new(move |socket, auth| {
                Box::pin(middleware(socket, auth))
            }));
    }

    pub async fn on<F, Fut>(&self, event: impl Into<String>, handler: F)
    where
        F: Fn(Socket, EventPayload) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = ()> + Send + 'static,
    {
        let state = self.io.ensure_namespace(&self.name).await;
        state.event_handlers.write().await.insert(
            event.into(),
            Arc::new(move |socket, payload| Box::pin(handler(socket, payload))),
        );
    }

    pub async fn emit<T: Serialize>(&self, event: &str, data: T) -> Result<()> {
        self.to_all().emit(event, data).await
    }

    pub fn to(&self, room: impl Into<String>) -> Broadcast {
        let mut rooms = HashSet::new();
        rooms.insert(room.into());
        Broadcast {
            io: self.io.clone(),
            opts: BroadcastOptions {
                namespace: self.name.clone(),
                rooms,
                ..BroadcastOptions::default()
            },
        }
    }

    pub fn to_all(&self) -> Broadcast {
        Broadcast {
            io: self.io.clone(),
            opts: BroadcastOptions {
                namespace: self.name.clone(),
                ..BroadcastOptions::default()
            },
        }
    }
}

pub struct Broadcast {
    io: SocketIo,
    opts: BroadcastOptions,
}

impl Broadcast {
    pub fn to(mut self, room: impl Into<String>) -> Self {
        self.opts.rooms.insert(room.into());
        self
    }

    pub fn except(mut self, room: impl Into<String>) -> Self {
        self.opts.except.insert(room.into());
        self
    }

    pub async fn emit<T: Serialize>(self, event: &str, data: T) -> Result<()> {
        self.io.broadcast_with_opts(self.opts, event, data).await
    }
}

impl Default for NamespaceState {
    fn default() -> Self {
        Self {
            connect_handler: RwLock::new(None),
            disconnect_handler: RwLock::new(None),
            event_handlers: RwLock::new(HashMap::new()),
            middleware: RwLock::new(Vec::new()),
        }
    }
}