swactor/src/delivery.rs

use std::any::Any;
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, OnceLock, RwLock};
use std::thread::Thread;

use crate::actor::{ActorAddress, AnyActor, Message, MonitorRef};
use crate::channel::Sender;
use crate::config::RuntimeConfig;
use crate::stats::WorkerStats;
use crate::Error;

// ─── Address Map Types ───────────────────────────────────────────────────────

/// Identifies a worker thread.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct WorkerId(pub(crate) usize);

impl WorkerId {
    pub fn as_usize(self) -> usize {
        self.0
    }
}

/// Maps actor addresses to the worker that owns them.
///
/// `RwLock<HashMap>` — zero contention for parallel reads, write-rare (only on spawn).
pub(crate) struct AddressMap {
    inner: RwLock<HashMap<ActorAddress, WorkerId>>,
}

impl AddressMap {
    pub fn with_capacity(cap: usize) -> Self {
        Self {
            inner: RwLock::new(HashMap::with_capacity(cap)),
        }
    }

    pub fn insert(&self, addr: ActorAddress, worker: WorkerId) {
        self.inner.write().unwrap().insert(addr, worker);
    }

    pub fn lookup(&self, addr: &ActorAddress) -> Option<WorkerId> {
        self.inner.read().unwrap().get(addr).copied()
    }

    /// Remove an actor address from the map (e.g., after permanent poisoning).
    pub fn remove(&self, addr: &ActorAddress) {
        self.inner.write().unwrap().remove(addr);
    }

    /// Returns a snapshot of all (address, worker) pairs.
    pub fn snapshot(&self) -> Vec<(ActorAddress, WorkerId)> {
        self.inner
            .read()
            .unwrap()
            .iter()
            .map(|(addr, wid)| (*addr, *wid))
            .collect()
    }
}

/// Load-aware actor placement strategy.
///
/// Picks the worker with the lowest load score (actor count + mailbox depth).
/// When all workers have equal load (e.g., before any ticks), falls back to
/// round-robin via a rotating start position for the scan.
pub(crate) struct Placement {
    next: AtomicUsize,
    num_workers: usize,
    worker_stats: Vec<Arc<WorkerStats>>,
}

impl Placement {
    pub fn new(num_workers: usize, worker_stats: Vec<Arc<WorkerStats>>) -> Self {
        Self {
            next: AtomicUsize::new(0),
            num_workers,
            worker_stats,
        }
    }

    pub fn next_worker(&self) -> WorkerId {
        let n = self.num_workers;
        if n == 1 {
            return WorkerId(0);
        }

        // Rotate the scan start for round-robin tie-breaking
        let rr = self.next.fetch_add(1, Ordering::Relaxed);

        let mut best_id = rr % n;
        let mut best_score = usize::MAX;

        for offset in 0..n {
            let i = (rr + offset) % n;
            let actors = self.worker_stats[i].num_actors.load(Ordering::Relaxed);
            let depth = self.worker_stats[i].total_mailbox_depth.load(Ordering::Relaxed);
            let score = actors + depth;
            if score < best_score {
                best_score = score;
                best_id = i;
            }
        }

        WorkerId(best_id)
    }
}

// ─── Delivery Types ──────────────────────────────────────────────────────────

/// A type-erased message envelope for cross-worker delivery.
///
/// Uses `Box` (no atomic refcount) and move semantics (no clone).
pub(crate) struct Envelope {
    dest: ActorAddress,
    payload: Box<dyn Any + Send>,
}

impl Envelope {
    pub fn new(dest: ActorAddress, payload: Box<dyn Any + Send>) -> Self {
        Self { dest, payload }
    }

    pub fn dest(&self) -> ActorAddress {
        self.dest
    }

    pub fn into_payload(self) -> Box<dyn Any + Send> {
        self.payload
    }
}

/// Type-erased sender for external inboxes.
pub(crate) trait SenderT: Send + Sync {
    fn try_send_any(&self, msg: Box<dyn Any + Send>);
}

impl<M: Message> SenderT for Sender<M> {
    fn try_send_any(&self, msg: Box<dyn Any + Send>) {
        if let Ok(typed) = msg.downcast::<M>() {
            Sender::send(self, *typed);
        }
    }
}

/// Registry of external inboxes — replaces the Router's role for non-actor receivers.
pub(crate) struct InboxRegistry {
    senders: RwLock<HashMap<ActorAddress, Arc<dyn SenderT>>>,
}

impl InboxRegistry {
    pub fn new() -> Self {
        Self {
            senders: RwLock::new(HashMap::new()),
        }
    }

    pub fn register(&self, addr: ActorAddress, sender: Arc<dyn SenderT>) {
        self.senders.write().unwrap().insert(addr, sender);
    }

    /// Check if an address is registered without consuming a message.
    #[cfg(feature = "transport")]
    pub fn contains(&self, addr: &ActorAddress) -> bool {
        self.senders.read().unwrap().contains_key(addr)
    }

    pub fn try_deliver(
        &self,
        addr: ActorAddress,
        msg: Box<dyn Any + Send>,
    ) -> Result<(), Error> {
        let senders = self.senders.read().unwrap();
        if let Some(sender) = senders.get(&addr) {
            sender.try_send_any(msg);
            Ok(())
        } else {
            Err(Error::from("Address not found"))
        }
    }
}

/// Shared state passed to tick_once — single thin pointer avoids register spill.
pub(crate) struct TickContext<'a> {
    pub(crate) address_map: &'a AddressMap,
    pub(crate) transfer_txs: &'a [Sender<Envelope>],
    pub(crate) spawn_txs: &'a [Sender<(ActorAddress, Box<dyn AnyActor>)>],
    pub(crate) placement: &'a Placement,
    pub(crate) inbox_registry: &'a InboxRegistry,
    pub(crate) config: &'a RuntimeConfig,
    pub(crate) name_registry: &'a NameRegistry,
    pub(crate) monitor_registry: &'a MonitorRegistry,
    pub(crate) group_registry: &'a GroupRegistry,
    pub(crate) stats_hook: Option<&'a dyn crate::stats::StatsHook>,
    /// Thread handles for waking parked workers on cross-worker sends.
    pub(crate) worker_threads: &'a [OnceLock<Thread>],
    #[cfg(feature = "transport")]
    pub(crate) codec_registry: Option<&'a crate::transport::CodecRegistry>,
    #[cfg(feature = "transport")]
    pub(crate) transport_router: Option<&'a crate::transport::TransportRouter>,
}

// ─── Name Registry ──────────────────────────────────────────────────────────

/// Named actor registry — maps human-readable names to actor addresses.
///
/// `RwLock<HashMap>` — same pattern as `AddressMap`. Write-rare (spawn/death),
/// read-often (lookup). A reverse map enables O(1) cleanup on actor death.
pub(crate) struct NameRegistry {
    names: RwLock<HashMap<String, ActorAddress>>,
    reverse: RwLock<HashMap<ActorAddress, String>>,
}

impl NameRegistry {
    pub fn new() -> Self {
        Self {
            names: RwLock::new(HashMap::new()),
            reverse: RwLock::new(HashMap::new()),
        }
    }

    /// Register a name → address mapping. Returns `Err` if the name is already taken.
    pub fn register(&self, name: String, addr: ActorAddress) -> Result<(), crate::Error> {
        let mut names = self.names.write().unwrap();
        if names.contains_key(&name) {
            return Err(crate::Error::from("Name already registered"));
        }
        names.insert(name.clone(), addr);
        drop(names);
        self.reverse.write().unwrap().insert(addr, name);
        Ok(())
    }

    /// Look up an actor address by name.
    pub fn lookup(&self, name: &str) -> Option<ActorAddress> {
        self.names.read().unwrap().get(name).copied()
    }

    /// Unregister a name, returning the address it was bound to.
    pub fn unregister(&self, name: &str) -> Option<ActorAddress> {
        let addr = self.names.write().unwrap().remove(name)?;
        self.reverse.write().unwrap().remove(&addr);
        Some(addr)
    }

    /// Remove a name by address (called on actor death for auto-cleanup).
    pub fn unregister_by_addr(&self, addr: &ActorAddress) {
        if let Some(name) = self.reverse.write().unwrap().remove(addr) {
            self.names.write().unwrap().remove(&name);
        }
    }

    /// Return all registered names.
    pub fn registered_names(&self) -> Vec<String> {
        self.names.read().unwrap().keys().cloned().collect()
    }
}

// ─── Monitor Registry ────────────────────────────────────────────────────────

/// Tracks monitor subscriptions: watched actor → list of (MonitorRef, watcher address).
///
/// Write-rare (monitor/demonitor/death), read at cleanup time.
pub(crate) struct MonitorRegistry {
    /// watched_addr → [(mref, watcher_addr)]
    monitors: RwLock<HashMap<ActorAddress, Vec<(MonitorRef, ActorAddress)>>>,
    /// mref → watched_addr (for O(1) demonitor)
    ref_to_target: RwLock<HashMap<MonitorRef, ActorAddress>>,
    next_ref: AtomicU64,
}

impl MonitorRegistry {
    pub fn new() -> Self {
        Self {
            monitors: RwLock::new(HashMap::new()),
            ref_to_target: RwLock::new(HashMap::new()),
            next_ref: AtomicU64::new(1),
        }
    }

    /// Register a monitor: `watcher` wants to know when `target` dies.
    pub fn register(&self, watcher: ActorAddress, target: ActorAddress) -> MonitorRef {
        let id = self.next_ref.fetch_add(1, Ordering::Relaxed);
        let mref = MonitorRef(id);
        self.monitors.write().unwrap()
            .entry(target)
            .or_default()
            .push((mref, watcher));
        self.ref_to_target.write().unwrap().insert(mref, target);
        mref
    }

    /// Cancel a monitor by its ref.
    pub fn deregister(&self, mref: MonitorRef) {
        if let Some(target) = self.ref_to_target.write().unwrap().remove(&mref) {
            let mut monitors = self.monitors.write().unwrap();
            if let Some(watchers) = monitors.get_mut(&target) {
                watchers.retain(|(r, _)| *r != mref);
                if watchers.is_empty() {
                    monitors.remove(&target);
                }
            }
        }
    }

    /// Remove and return all monitors for a dead actor.
    pub fn take_monitors(&self, target: &ActorAddress) -> Vec<(MonitorRef, ActorAddress)> {
        let watchers = self.monitors.write().unwrap().remove(target).unwrap_or_default();
        let mut ref_map = self.ref_to_target.write().unwrap();
        for (mref, _) in &watchers {
            ref_map.remove(mref);
        }
        watchers
    }

    /// Remove all monitor subscriptions where `addr` is the watcher (dead watcher cleanup).
    pub fn remove_watcher(&self, addr: &ActorAddress) {
        let mut monitors = self.monitors.write().unwrap();
        let mut ref_map = self.ref_to_target.write().unwrap();
        // Iterate all targets and remove entries where this addr is the watcher
        monitors.retain(|_target, watchers| {
            watchers.retain(|(mref, watcher)| {
                if watcher == addr {
                    ref_map.remove(mref);
                    false
                } else {
                    true
                }
            });
            !watchers.is_empty()
        });
    }
}

// ─── Group Registry ─────────────────────────────────────────────────────────

/// Actor groups (pub-sub). Actors join/leave named groups; messages can be
/// broadcast to all members of a group.
///
/// Groups are created lazily on first join and removed when empty.
pub(crate) struct GroupRegistry {
    /// group_name → set of member addresses
    groups: RwLock<HashMap<String, HashSet<ActorAddress>>>,
    /// actor_addr → set of group names (reverse map for O(G) cleanup on death)
    memberships: RwLock<HashMap<ActorAddress, HashSet<String>>>,
}

impl GroupRegistry {
    pub fn new() -> Self {
        Self {
            groups: RwLock::new(HashMap::new()),
            memberships: RwLock::new(HashMap::new()),
        }
    }

    /// Add an actor to a named group. Group is created if it doesn't exist.
    pub fn join(&self, group: String, addr: ActorAddress) {
        self.groups.write().unwrap()
            .entry(group.clone())
            .or_default()
            .insert(addr);
        self.memberships.write().unwrap()
            .entry(addr)
            .or_default()
            .insert(group);
    }

    /// Remove an actor from a named group. Empty groups are auto-deleted.
    pub fn leave(&self, group: &str, addr: &ActorAddress) {
        let mut groups = self.groups.write().unwrap();
        if let Some(members) = groups.get_mut(group) {
            members.remove(addr);
            if members.is_empty() {
                groups.remove(group);
            }
        }
        drop(groups);
        if let Some(membership) = self.memberships.write().unwrap().get_mut(addr) {
            membership.remove(group);
        }
    }

    /// Return all members of a group.
    pub fn members(&self, group: &str) -> Vec<ActorAddress> {
        self.groups.read().unwrap()
            .get(group)
            .map(|s| s.iter().copied().collect())
            .unwrap_or_default()
    }

    /// Remove a dead actor from all its groups.
    pub fn cleanup(&self, addr: &ActorAddress) {
        let group_names = self.memberships.write().unwrap().remove(addr);
        if let Some(names) = group_names {
            let mut groups = self.groups.write().unwrap();
            for name in names {
                if let Some(members) = groups.get_mut(&name) {
                    members.remove(addr);
                    if members.is_empty() {
                        groups.remove(&name);
                    }
                }
            }
        }
    }

    /// Return all active group names.
    pub fn group_names(&self) -> Vec<String> {
        self.groups.read().unwrap().keys().cloned().collect()
    }
}

impl<'a> TickContext<'a> {
    /// Route a message whose destination is not in the local address map.
    /// Tries inbox registry, then remote transport, then falls back to inbox error.
    pub(crate) fn route_nonlocal(
        &self,
        addr: ActorAddress,
        msg: Box<dyn Any + Send>,
    ) -> Result<(), Error> {
        #[cfg(feature = "transport")]
        {
            if self.inbox_registry.contains(&addr) {
                return self.inbox_registry.try_deliver(addr, msg);
            }
            if let (Some(cr), Some(tr)) = (self.codec_registry, self.transport_router) {
                return crate::transport::send_via_transport(addr, msg, cr, tr);
            }
        }
        self.inbox_registry.try_deliver(addr, msg)
    }
}