imks/socket/namespace.rs

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

use dashmap::DashMap;
use tokio::sync::RwLock;

use crate::socket::adapter::{Adapter, BroadcastFlags, BroadcastOptions};
use crate::socket::packet::Packet;
use crate::socket::socket::Socket;

pub type EventHandler = Arc<dyn Fn(Arc<Socket>, &serde_json::Value) + Send + Sync>;
type ConnectHandler =
    Arc<dyn Fn(&Socket, Option<&serde_json::Value>) -> Result<(), String> + Send + Sync>;

pub struct Namespace {
    pub path: String,
    /// Primary storage: socket_sid → Socket
    sockets: DashMap<String, Arc<Socket>>,
    /// Reverse index: engine_sid → socket_sid (for engine-level lookups)
    engine_to_socket: DashMap<String, String>,
    handlers: RwLock<HashMap<String, Vec<EventHandler>>>,
    connect_handler: RwLock<Option<ConnectHandler>>,
    rooms: DashMap<String, HashSet<String>>,
    socket_rooms: DashMap<String, HashSet<String>>,
    pub(crate) adapter: RwLock<Option<Arc<dyn Adapter>>>,
}

impl Namespace {
    pub fn new(path: impl Into<String>) -> Self {
        Self {
            path: path.into(),
            sockets: DashMap::new(),
            engine_to_socket: DashMap::new(),
            handlers: RwLock::new(HashMap::new()),
            connect_handler: RwLock::new(None),
            rooms: DashMap::new(),
            socket_rooms: DashMap::new(),
            adapter: RwLock::new(None),
        }
    }

    pub async fn set_adapter(&self, adapter: Arc<dyn Adapter>) {
        let mut guard = self.adapter.write().await;
        *guard = Some(adapter);
    }

    /// Add a socket to this namespace. Returns Err if the connect handler rejects.
    pub async fn add_socket(
        &self,
        socket: Arc<Socket>,
        auth_data: Option<&serde_json::Value>,
    ) -> Result<(), String> {
        // Run connect handler before adding to storage
        let handler = self.connect_handler.read().await;
        if let Some(ref h) = *handler {
            h(&socket, auth_data)?;
        }
        drop(handler);

        let socket_sid = socket.sid.clone();
        let engine_sid = socket.engine_sid.clone();

        // Register with adapter (socket_sid → engine_sid mapping)
        let adapter = self.adapter.read().await;
        if let Some(ref adapter) = *adapter
            && let Err(e) = adapter.register(&socket_sid, &engine_sid, &self.path).await
        {
            tracing::warn!("Adapter register error for socket {}: {}", socket_sid, e);
        }

        // Store socket by socket_sid, plus reverse index
        self.sockets.insert(socket_sid.clone(), socket);
        self.engine_to_socket.insert(engine_sid, socket_sid);
        Ok(())
    }

    /// Remove a socket by its socket SID.
    pub async fn remove_socket_by_sid(&self, socket_sid: &str) {
        if let Some((_, socket)) = self.sockets.remove(socket_sid) {
            self.engine_to_socket.remove(&socket.engine_sid);
            self.remove_socket_from_local_rooms(socket_sid);

            let adapter = self.adapter.read().await;
            if let Some(ref adapter) = *adapter
                && let Err(e) = adapter.del_all(socket_sid, &self.path).await
            {
                tracing::warn!("Adapter del_all error for socket {}: {}", socket_sid, e);
            }
        }
    }

    /// Remove a socket by its engine SID (for engine-level disconnections).
    pub async fn remove_socket(&self, engine_sid: &str) {
        if let Some((_, socket_sid)) = self.engine_to_socket.remove(engine_sid) {
            self.remove_socket_by_sid(&socket_sid).await;
        }
    }

    /// Look up a socket by its socket SID.
    pub fn get_socket(&self, socket_sid: &str) -> Option<Arc<Socket>> {
        self.sockets.get(socket_sid).map(|r| r.value().clone())
    }

    /// Look up a socket by its engine SID (reverse lookup).
    pub fn get_socket_by_engine_sid(&self, engine_sid: &str) -> Option<Arc<Socket>> {
        self.engine_to_socket
            .get(engine_sid)
            .and_then(|entry| self.sockets.get(entry.value()).map(|r| r.value().clone()))
    }

    pub fn socket_count(&self) -> usize {
        self.sockets.len()
    }

    pub async fn on_event(&self, event: impl Into<String>, handler: EventHandler) {
        let mut handlers = self.handlers.write().await;
        handlers.entry(event.into()).or_default().push(handler);
    }

    pub async fn on_connect<F>(&self, handler: F)
    where
        F: Fn(&Socket, Option<&serde_json::Value>) -> Result<(), String> + Send + Sync + 'static,
    {
        let mut connect_handler = self.connect_handler.write().await;
        *connect_handler = Some(Arc::new(handler));
    }

    pub async fn emit(&self, event: impl Into<String>, data: serde_json::Value) {
        let event_name = event.into();
        let packet = Packet::event(&self.path, serde_json::json!([event_name, data]), None);

        let adapter = self.adapter.read().await;
        if let Some(ref adapter) = *adapter {
            let opts = BroadcastOptions::default();
            if let Err(e) = adapter.broadcast(&packet, &opts).await {
                tracing::warn!("Adapter broadcast error: {}", e);
            }
        } else {
            self.emit_local(&packet);
        }
    }

    pub async fn emit_to_room(
        &self,
        room: &str,
        event: impl Into<String>,
        data: serde_json::Value,
    ) {
        let event_name = event.into();
        let packet = Packet::event(&self.path, serde_json::json!([event_name, data]), None);

        let adapter = self.adapter.read().await;
        if let Some(ref adapter) = *adapter {
            let opts = BroadcastOptions {
                rooms: HashSet::from([room.to_string()]),
                except: HashSet::new(),
                flags: BroadcastFlags::default(),
            };
            if let Err(e) = adapter.broadcast(&packet, &opts).await {
                tracing::warn!("Adapter broadcast to room error: {}", e);
            }
        } else {
            self.emit_local_to_room(&packet, room, &HashSet::new());
        }
    }

    pub fn emit_local(&self, packet: &Packet) {
        for entry in self.sockets.iter() {
            self.send_local_packet(entry.value(), packet);
        }
    }

    pub fn emit_local_filtered(&self, packet: &Packet, opts: &BroadcastOptions) {
        if opts.rooms.is_empty() {
            for entry in self.sockets.iter() {
                if !opts.except.contains(entry.key()) {
                    self.send_local_packet(entry.value(), packet);
                }
            }
            return;
        }

        let mut target_sids = HashSet::new();
        for room in &opts.rooms {
            if let Some(room_sids) = self.rooms.get(room) {
                target_sids.extend(room_sids.value().iter().cloned());
            }
        }

        for sid in target_sids {
            if opts.except.contains(&sid) {
                continue;
            }
            if let Some(socket) = self.get_socket(&sid) {
                self.send_local_packet(&socket, packet);
            }
        }
    }

    fn emit_local_to_room(&self, packet: &Packet, room: &str, except: &HashSet<String>) {
        let opts = BroadcastOptions {
            rooms: HashSet::from([room.to_string()]),
            except: except.clone(),
            flags: BroadcastFlags::default(),
        };
        self.emit_local_filtered(packet, &opts);
    }

    fn send_local_packet(&self, socket: &Socket, packet: &Packet) {
        if socket.send_packet(packet).is_err() {
            tracing::warn!("Failed to send event to socket {}", socket.sid);
        }
    }

    pub async fn emit_to(
        &self,
        socket_sid: &str,
        event: impl Into<String>,
        data: serde_json::Value,
    ) {
        if let Some(socket) = self.get_socket(socket_sid) {
            let event_name = event.into();
            let packet = Packet::event(&self.path, serde_json::json!([event_name, data]), None);
            if socket.send_packet(&packet).is_err() {
                tracing::warn!("Failed to send event to socket {}", socket.sid);
            }
        }
    }

    pub async fn handle_event(&self, socket: Arc<Socket>, event: &str, data: &serde_json::Value) {
        let handlers = self.handlers.read().await;
        if let Some(event_handlers) = handlers.get(event) {
            for handler in event_handlers {
                handler(Arc::clone(&socket), data);
            }
        }
    }

    pub async fn join_room(&self, socket_sid: &str, room: &str) -> crate::ImksResult<()> {
        if !self.sockets.contains_key(socket_sid) {
            return Err(crate::ImksError::SocketNotFound(socket_sid.to_string()));
        }

        self.rooms
            .entry(room.to_string())
            .or_default()
            .value_mut()
            .insert(socket_sid.to_string());
        self.socket_rooms
            .entry(socket_sid.to_string())
            .or_default()
            .value_mut()
            .insert(room.to_string());

        let adapter = self.adapter.read().await;
        if let Some(ref adapter) = *adapter
            && let Err(e) = adapter.add(socket_sid, room, &self.path).await
        {
            self.remove_local_room(socket_sid, room);
            return Err(e.into());
        }
        Ok(())
    }

    pub async fn leave_room(&self, socket_sid: &str, room: &str) -> crate::ImksResult<()> {
        let adapter = self.adapter.read().await;
        if let Some(ref adapter) = *adapter {
            adapter.del(socket_sid, room, &self.path).await?;
        }

        self.remove_local_room(socket_sid, room);
        Ok(())
    }

    fn remove_local_room(&self, socket_sid: &str, room: &str) {
        if let Some(mut sids) = self.rooms.get_mut(room) {
            sids.value_mut().remove(socket_sid);
            if sids.value().is_empty() {
                drop(sids);
                self.rooms.remove(room);
            }
        }

        if let Some(mut rooms) = self.socket_rooms.get_mut(socket_sid) {
            rooms.value_mut().remove(room);
            if rooms.value().is_empty() {
                drop(rooms);
                self.socket_rooms.remove(socket_sid);
            }
        }
    }

    fn remove_socket_from_local_rooms(&self, socket_sid: &str) {
        if let Some((_, rooms)) = self.socket_rooms.remove(socket_sid) {
            for room in rooms {
                if let Some(mut sids) = self.rooms.get_mut(&room) {
                    sids.value_mut().remove(socket_sid);
                    if sids.value().is_empty() {
                        drop(sids);
                        self.rooms.remove(&room);
                    }
                }
            }
        }
    }
}

pub struct NamespaceManager {
    namespaces: DashMap<String, Arc<Namespace>>,
}

impl NamespaceManager {
    pub fn new() -> Self {
        let manager = Self {
            namespaces: DashMap::new(),
        };
        manager.create_namespace("/");
        manager
    }

    pub fn create_namespace(&self, path: impl Into<String>) -> Arc<Namespace> {
        let path = path.into();
        let namespace = Arc::new(Namespace::new(&path));
        self.namespaces.insert(path.clone(), namespace.clone());
        namespace
    }

    pub fn get_namespace(&self, path: &str) -> Option<Arc<Namespace>> {
        self.namespaces.get(path).map(|r| r.value().clone())
    }

    pub fn get_or_create_namespace(&self, path: &str) -> Arc<Namespace> {
        if let Some(ns) = self.get_namespace(path) {
            ns
        } else {
            self.create_namespace(path)
        }
    }

    pub fn remove_namespace(&self, path: &str) {
        self.namespaces.remove(path);
    }

    pub fn namespace_count(&self) -> usize {
        self.namespaces.len()
    }

    pub fn all_namespaces(&self) -> Vec<Arc<Namespace>> {
        self.namespaces.iter().map(|e| e.value().clone()).collect()
    }
}

impl Default for NamespaceManager {
    fn default() -> Self {
        Self::new()
    }
}

/// Validate a namespace path. Returns true if the path is valid.
/// Rules: must start with '/', max 256 chars, no control characters.
pub fn is_valid_namespace(path: &str) -> bool {
    !path.is_empty()
        && path.starts_with('/')
        && path.len() <= 256
        && !path.chars().any(|c| c.is_control())
}