imks/socket/adapter/local.rs

use std::collections::HashSet;
use std::sync::Arc;

use async_trait::async_trait;
use dashmap::DashMap;
use uuid::Uuid;

use crate::socket::adapter::{Adapter, AdapterError, BroadcastOptions, LocalSendFn, SocketInfo};
use crate::socket::packet::Packet;

pub struct LocalAdapter {
    server_id: String,
    rooms: Arc<DashMap<String, HashSet<String>>>,
    socket_rooms: Arc<DashMap<String, HashSet<String>>>,
    /// socket_sid → engine_sid
    pub socket_sids: Arc<DashMap<String, String>>,
    /// socket_sid → namespace path
    socket_namespace: Arc<DashMap<String, String>>,
    send_fn: LocalSendFn,
}

impl LocalAdapter {
    pub fn new(
        send_fn: impl Fn(&str, &Packet) -> Result<(), String> + Send + Sync + 'static,
    ) -> Self {
        Self {
            server_id: Uuid::new_v4().to_string(),
            rooms: Arc::new(DashMap::new()),
            socket_rooms: Arc::new(DashMap::new()),
            socket_sids: Arc::new(DashMap::new()),
            socket_namespace: Arc::new(DashMap::new()),
            send_fn: Arc::new(send_fn),
        }
    }

    fn room_key(ns: &str, room: &str) -> String {
        format!("{}:{}", ns, room)
    }

    /// Collect socket SIDs matching the broadcast options, scoped to the given namespace.
    fn collect_matching_sids(&self, opts: &BroadcastOptions, namespace: &str) -> Vec<String> {
        if opts.rooms.is_empty() {
            // Broadcast to all sockets in this namespace only
            self.socket_sids
                .iter()
                .filter(|e| {
                    self.socket_namespace
                        .get(e.key())
                        .map(|ns| ns.value() == namespace)
                        .unwrap_or(false)
                })
                .map(|e| e.key().clone())
                .collect()
        } else {
            let mut sids = HashSet::new();
            for room in &opts.rooms {
                let key = Self::room_key(namespace, room);
                if let Some(entry) = self.rooms.get(&key) {
                    for sid in entry.value() {
                        sids.insert(sid.clone());
                    }
                }
            }
            sids.into_iter().collect()
        }
    }
}

#[async_trait]
impl Adapter for LocalAdapter {
    async fn broadcast(
        &self,
        packet: &Packet,
        opts: &BroadcastOptions,
    ) -> Result<(), AdapterError> {
        let namespace = &packet.namespace;
        let sids = self.collect_matching_sids(opts, namespace);
        for sid in &sids {
            if opts.except.contains(sid) {
                continue;
            }
            // socket_sids maps socket SID -> engine SID
            if let Some(entry) = self.socket_sids.get(sid) {
                let engine_sid = entry.value();
                let result = (self.send_fn)(engine_sid, packet);
                if let Err(e) = result {
                    tracing::warn!("Failed to broadcast to {}: {}", sid, e);
                }
            }
        }
        Ok(())
    }

    async fn register(
        &self,
        socket_sid: &str,
        engine_sid: &str,
        ns: &str,
    ) -> Result<(), AdapterError> {
        self.socket_sids
            .insert(socket_sid.to_string(), engine_sid.to_string());
        self.socket_namespace
            .insert(socket_sid.to_string(), ns.to_string());
        Ok(())
    }

    async fn unregister(&self, socket_sid: &str, ns: &str) -> Result<(), AdapterError> {
        self.del_all(socket_sid, ns).await
    }

    async fn add(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError> {
        let key = Self::room_key(ns, room);
        self.rooms
            .entry(key)
            .or_default()
            .value_mut()
            .insert(sid.to_string());
        self.socket_rooms
            .entry(sid.to_string())
            .or_default()
            .value_mut()
            .insert(room.to_string());
        Ok(())
    }

    async fn del(&self, sid: &str, room: &str, ns: &str) -> Result<(), AdapterError> {
        let key = Self::room_key(ns, room);
        if let Some(mut room_sids) = self.rooms.get_mut(&key) {
            room_sids.value_mut().remove(sid);
            if room_sids.value_mut().is_empty() {
                drop(room_sids);
                self.rooms.remove(&key);
            }
        }
        if let Some(mut rooms) = self.socket_rooms.get_mut(sid) {
            rooms.value_mut().remove(room);
            if rooms.value_mut().is_empty() {
                drop(rooms);
                self.socket_rooms.remove(sid);
            }
        }
        Ok(())
    }

    async fn del_all(&self, sid: &str, ns: &str) -> Result<(), AdapterError> {
        if let Some((_, rooms)) = self.socket_rooms.remove(sid) {
            for room in &rooms {
                let key = Self::room_key(ns, room);
                if let Some(mut room_sids) = self.rooms.get_mut(&key) {
                    room_sids.value_mut().remove(sid);
                    if room_sids.value_mut().is_empty() {
                        drop(room_sids);
                        self.rooms.remove(&key);
                    }
                }
            }
        }
        self.socket_sids.remove(sid);
        self.socket_namespace.remove(sid);
        Ok(())
    }

    async fn fetch_sockets(
        &self,
        opts: &BroadcastOptions,
    ) -> Result<Vec<SocketInfo>, AdapterError> {
        // fetch_sockets needs namespace context; use an empty namespace to match all
        // (this method is typically called for inspection, not delivery)
        let sids: Vec<String> = if opts.rooms.is_empty() {
            self.socket_sids.iter().map(|e| e.key().clone()).collect()
        } else {
            let mut sids_set = HashSet::new();
            for room in &opts.rooms {
                for entry in self.rooms.iter() {
                    if entry.key().ends_with(&format!(":{}", room)) {
                        for sid in entry.value() {
                            sids_set.insert(sid.clone());
                        }
                    }
                }
            }
            sids_set.into_iter().collect()
        };
        let mut result = Vec::new();
        for sid in &sids {
            if opts.except.contains(sid) {
                continue;
            }
            if self.socket_sids.contains_key(sid) {
                let namespace = self
                    .socket_namespace
                    .get(sid)
                    .map(|r| r.value().clone())
                    .unwrap_or_default();
                let rooms = self
                    .socket_rooms
                    .get(sid)
                    .map(|r| r.value().clone())
                    .unwrap_or_default();
                result.push(SocketInfo {
                    sid: sid.clone(),
                    namespace,
                    rooms,
                });
            }
        }
        Ok(result)
    }

    async fn socket_rooms(&self, sid: &str) -> Result<HashSet<String>, AdapterError> {
        Ok(self
            .socket_rooms
            .get(sid)
            .map(|r| r.value().clone())
            .unwrap_or_default())
    }

    fn server_id(&self) -> &str {
        &self.server_id
    }

    async fn close(&self) -> Result<(), AdapterError> {
        Ok(())
    }
}