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, BroadcastOptions, BroadcastFlags};
use crate::socket::packet::Packet;
use crate::socket::socket::Socket;

pub type EventHandler = Arc<dyn Fn(&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>>,
    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),
            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>) -> Result<(), String> {
        // Run connect handler before adding to storage
        let handler = self.connect_handler.read().await;
        if let Some(ref h) = *handler {
            h(&socket, None)?;
        }
        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 {
            if 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);

            let adapter = self.adapter.read().await;
            if let Some(ref adapter) = *adapter {
                if 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(&packet);
        }
    }

    pub fn emit_local(&self, packet: &Packet) {
        for entry in self.sockets.iter() {
            let socket = entry.value();
            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: &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(socket, data);
            }
        }
    }
}

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())
}