ontoref-code/crates/ontoref-daemon/src/actors.rs

use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use tracing::{debug, info, warn};

/// Actor token: `{actor_type}:{hostname}:{pid}:{nonce}`.
///
/// The nanosecond nonce prevents PID-recycling impersonation attacks and
/// makes tokens unpredictable without NS-precision timing access. Tokens are
/// no longer idempotent across restarts — clients store their token from the
/// registration response and re-register after a daemon restart.
pub fn make_token(actor_type: &str, hostname: &str, pid: u32) -> String {
    let nonce: u32 = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.subsec_nanos())
        .unwrap_or_else(|_| pid.wrapping_mul(0x9e37_79b9));
    format!("{actor_type}:{hostname}:{pid}:{nonce:08x}")
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ActorSession {
    pub actor_type: String,
    pub hostname: String,
    pub pid: u32,
    pub project: String,
    /// Role identifier — validated against `.ontoref/roles.ncl` if present;
    /// defaults to "developer".
    pub role: String,
    pub registered_at: u64,
    pub last_seen: u64,
    pub pending_notifications: AtomicU64,
    /// UI and behavioral preferences (theme, nav_mode, etc.).
    pub preferences: serde_json::Value,
}

/// Serialisable snapshot of an `ActorSession` for disk persistence.
/// Written to `{persist_dir}/{token_safe}.json` on profile update.
#[derive(Debug, Serialize, Deserialize)]
struct PersistedSession {
    token: String,
    actor_type: String,
    hostname: String,
    pid: u32,
    project: String,
    role: String,
    registered_at: u64,
    last_seen: u64,
    preferences: serde_json::Value,
}

/// Lock-free concurrent actor registry.
///
/// Key = token (`actor_type:hostname:pid`), value = session metadata.
/// Actors register via HTTP POST and deregister via HTTP DELETE or
/// are reaped by the periodic sweep.
///
/// When `persist_dir` is set (e.g. `{project_root}/.ontoref/sessions`), actor
/// profiles are written as JSON files on update and loaded on construction.
pub struct ActorRegistry {
    sessions: DashMap<String, ActorSession>,
    stale_timeout_secs: u64,
    /// Cached at construction to avoid syscall on every sweep iteration.
    local_hostname: String,
    /// Directory for session profile persistence (`{root}/.ontoref/sessions`).
    persist_dir: Option<PathBuf>,
}

impl ActorRegistry {
    pub fn new(stale_timeout_secs: u64) -> Self {
        let local_hostname = hostname::get()
            .map(|h| h.to_string_lossy().to_string())
            .unwrap_or_default();
        Self {
            sessions: DashMap::new(),
            stale_timeout_secs,
            local_hostname,
            persist_dir: None,
        }
    }

    pub fn with_persist_dir(mut self, dir: PathBuf) -> Self {
        self.persist_dir = Some(dir);
        self
    }

    /// Load persisted session profiles from `persist_dir`. Tokens that do not
    /// correspond to a live actor are kept in the registry with their last-seen
    /// timestamp so the daemon can restore preferences on the next registration
    /// of the same token within the stale window.
    pub fn load_persisted(&self) {
        let Some(dir) = &self.persist_dir else { return };
        let Ok(entries) = std::fs::read_dir(dir) else {
            return;
        };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) != Some("json") {
                continue;
            }
            let Ok(bytes) = std::fs::read(&path) else {
                continue;
            };
            let Ok(session) = serde_json::from_slice::<PersistedSession>(&bytes) else {
                warn!(path = %path.display(), "failed to deserialise persisted session");
                continue;
            };
            // Do not re-insert if already registered (live session takes priority).
            if self.sessions.contains_key(&session.token) {
                continue;
            }
            let token = session.token.clone();
            self.sessions.insert(
                token.clone(),
                ActorSession {
                    actor_type: session.actor_type,
                    hostname: session.hostname,
                    pid: session.pid,
                    project: session.project,
                    role: session.role,
                    registered_at: session.registered_at,
                    last_seen: session.last_seen,
                    pending_notifications: AtomicU64::new(0),
                    preferences: session.preferences,
                },
            );
            debug!(token = %token, "loaded persisted actor session");
        }
    }

    pub fn register(&self, req: RegisterRequest) -> String {
        let token = make_token(&req.actor_type, &req.hostname, req.pid);
        let now = epoch_secs();

        let session = ActorSession {
            actor_type: req.actor_type,
            hostname: req.hostname,
            pid: req.pid,
            project: req.project,
            role: req.role.unwrap_or_else(|| "developer".to_string()),
            registered_at: now,
            last_seen: now,
            pending_notifications: AtomicU64::new(0),
            preferences: req
                .preferences
                .unwrap_or(serde_json::Value::Object(Default::default())),
        };

        self.sessions.insert(token.clone(), session);
        info!(token = %token, "actor registered");
        token
    }

    pub fn deregister(&self, token: &str) -> bool {
        let removed = self.sessions.remove(token).is_some();
        if removed {
            info!(token = %token, "actor deregistered");
        }
        removed
    }

    /// Update role and/or preferences for a registered session. Returns `true`
    /// if found. Persists the updated profile to disk if `persist_dir` is
    /// configured.
    pub fn update_profile(
        &self,
        token: &str,
        role: Option<String>,
        preferences: Option<serde_json::Value>,
    ) -> bool {
        let Some(mut session) = self.sessions.get_mut(token) else {
            return false;
        };
        if let Some(r) = role {
            session.role = r;
        }
        if let Some(p) = preferences {
            session.preferences = p;
        }
        session.last_seen = epoch_secs();

        if let Some(dir) = &self.persist_dir {
            let persisted = PersistedSession {
                token: token.to_string(),
                actor_type: session.actor_type.clone(),
                hostname: session.hostname.clone(),
                pid: session.pid,
                project: session.project.clone(),
                role: session.role.clone(),
                registered_at: session.registered_at,
                last_seen: session.last_seen,
                preferences: session.preferences.clone(),
            };
            drop(session); // release DashMap write guard before I/O
            self.write_persisted_session(dir, token, &persisted);
        }
        true
    }

    fn write_persisted_session(&self, dir: &PathBuf, token: &str, session: &PersistedSession) {
        if let Err(e) = std::fs::create_dir_all(dir) {
            warn!(dir = %dir.display(), error = %e, "failed to create sessions persist dir");
            return;
        }
        let file = dir.join(format!("{}.json", token_to_filename(token)));
        match serde_json::to_vec_pretty(session) {
            Ok(bytes) => {
                if let Err(e) = std::fs::write(&file, bytes) {
                    warn!(path = %file.display(), error = %e, "failed to persist actor session");
                }
            }
            Err(e) => warn!(error = %e, "failed to serialise actor session"),
        }
    }

    /// Update `last_seen` timestamp. Returns `true` if the token was found.
    pub fn touch(&self, token: &str) -> bool {
        if let Some(mut session) = self.sessions.get_mut(token) {
            session.last_seen = epoch_secs();
            true
        } else {
            false
        }
    }

    pub fn increment_pending(&self, token: &str) {
        if let Some(session) = self.sessions.get(token) {
            session
                .pending_notifications
                .fetch_add(1, Ordering::Relaxed);
        }
    }

    pub fn clear_pending(&self, token: &str) {
        if let Some(session) = self.sessions.get(token) {
            session.pending_notifications.store(0, Ordering::Relaxed);
        }
    }

    pub fn get(&self, token: &str) -> Option<ActorSessionView> {
        self.sessions
            .get(token)
            .map(|s| ActorSessionView::from(&*s))
    }

    /// Atomically snapshot the session AND update `last_seen` under the same
    /// shard write lock. Eliminates the TOCTOU gap between a separate `get`
    /// call followed by `touch` — a sweep cannot remove the entry between them.
    pub fn get_and_touch(&self, token: &str) -> Option<ActorSessionView> {
        self.sessions.get_mut(token).map(|mut s| {
            s.last_seen = epoch_secs();
            ActorSessionView::from(&*s)
        })
    }

    pub fn list(&self) -> Vec<(String, ActorSessionView)> {
        self.sessions
            .iter()
            .map(|entry| (entry.key().clone(), ActorSessionView::from(entry.value())))
            .collect()
    }

    pub fn list_for_project(&self, project: &str) -> Vec<(String, ActorSessionView)> {
        self.sessions
            .iter()
            .filter(|entry| entry.value().project == project)
            .map(|entry| (entry.key().clone(), ActorSessionView::from(entry.value())))
            .collect()
    }

    /// Tokens of all actors registered to a given project.
    pub fn tokens_for_project(&self, project: &str) -> Vec<String> {
        self.sessions
            .iter()
            .filter(|entry| entry.value().project == project)
            .map(|entry| entry.key().clone())
            .collect()
    }

    /// Deregister all sessions belonging to a project. Returns the count
    /// removed. Used when a project's keys are rotated to force
    /// re-authentication.
    pub fn deregister_all_for_project(&self, project: &str) -> usize {
        let tokens: Vec<String> = self
            .sessions
            .iter()
            .filter(|entry| entry.value().project == project)
            .map(|entry| entry.key().clone())
            .collect();
        let count = tokens.len();
        for token in tokens {
            self.sessions.remove(&token);
        }
        count
    }

    /// Sweep stale sessions. Local actors: `kill -0` liveness check.
    /// Remote actors (CI): `last_seen` timeout.
    /// Returns tokens of reaped sessions.
    pub fn sweep_stale(&self) -> Vec<String> {
        let now = epoch_secs();
        let mut reaped = Vec::new();

        // Collect tokens to remove (can't remove during iteration with DashMap).
        let stale_tokens: Vec<String> = self
            .sessions
            .iter()
            .filter(|entry| {
                let session = entry.value();
                if session.hostname == self.local_hostname {
                    !process_alive(session.pid)
                } else {
                    now.saturating_sub(session.last_seen) > self.stale_timeout_secs
                }
            })
            .map(|entry| entry.key().clone())
            .collect();

        for token in stale_tokens {
            if self.sessions.remove(&token).is_some() {
                debug!(token = %token, "reaped stale actor session");
                if let Some(dir) = &self.persist_dir {
                    let _ = std::fs::remove_file(
                        dir.join(format!("{}.json", token_to_filename(&token))),
                    );
                }
                reaped.push(token);
            }
        }

        if !reaped.is_empty() {
            info!(count = reaped.len(), "sweep completed");
        }

        reaped
    }

    pub fn count(&self) -> usize {
        self.sessions.len()
    }

    /// Collect all unique project names from registered actors.
    pub fn active_projects(&self) -> HashSet<String> {
        self.sessions
            .iter()
            .map(|entry| entry.value().project.clone())
            .collect()
    }
}

/// Serializable view of an `ActorSession` (AtomicU64 → u64).
#[derive(Debug, Clone, Serialize)]
pub struct ActorSessionView {
    pub actor_type: String,
    pub hostname: String,
    pub pid: u32,
    pub project: String,
    pub role: String,
    pub registered_at: u64,
    pub last_seen: u64,
    pub pending_notifications: u64,
    pub preferences: serde_json::Value,
}

impl From<&ActorSession> for ActorSessionView {
    fn from(s: &ActorSession) -> Self {
        Self {
            actor_type: s.actor_type.clone(),
            hostname: s.hostname.clone(),
            pid: s.pid,
            project: s.project.clone(),
            role: s.role.clone(),
            registered_at: s.registered_at,
            last_seen: s.last_seen,
            pending_notifications: s.pending_notifications.load(Ordering::Relaxed),
            preferences: s.preferences.clone(),
        }
    }
}

#[derive(Debug, Default, Deserialize)]
pub struct RegisterRequest {
    pub actor_type: String,
    pub hostname: String,
    pub pid: u32,
    pub project: String,
    #[serde(default)]
    pub role: Option<String>,
    #[serde(default)]
    pub preferences: Option<serde_json::Value>,
}

/// Sanitise a token string into a safe filename component.
/// Replaces `:` and path-separator characters with `_`.
fn token_to_filename(token: &str) -> String {
    token
        .chars()
        .map(|c| {
            if c.is_alphanumeric() || c == '-' {
                c
            } else {
                '_'
            }
        })
        .collect()
}

fn epoch_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// Check if a local process is alive via `kill -0`.
#[cfg(unix)]
fn process_alive(pid: u32) -> bool {
    // SAFETY: kill(pid, 0) is a POSIX-standard liveness check.
    // Signal 0 is explicitly defined as "no signal sent, but error checking
    // is still performed" — it cannot affect the target process.
    unsafe { libc::kill(pid as libc::pid_t, 0) == 0 }
}

#[cfg(not(unix))]
fn process_alive(_pid: u32) -> bool {
    // Non-Unix: assume alive, rely on last_seen timeout.
    true
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn token_format() {
        // Token now includes an 8-hex-digit nanosecond nonce for entropy.
        let tok = make_token("developer", "macbook", 1234);
        let parts: Vec<&str> = tok.split(':').collect();
        assert_eq!(parts.len(), 4, "token must have 4 colon-separated parts");
        assert_eq!(parts[0], "developer");
        assert_eq!(parts[1], "macbook");
        assert_eq!(parts[2], "1234");
        assert_eq!(parts[3].len(), 8, "nonce must be 8 hex chars");
        assert!(
            parts[3].chars().all(|c| c.is_ascii_hexdigit()),
            "nonce must be hex"
        );
        // Two calls must produce different nonces (collisions possible but
        // vanishingly unlikely — ns-precision timestamps should differ).
        let tok2 = make_token("developer", "macbook", 1234);
        // Both are valid format; we don't assert inequality in case of same ns.
        let p2: Vec<&str> = tok2.split(':').collect();
        assert_eq!(p2.len(), 4);
    }

    #[test]
    fn register_and_deregister() {
        let registry = ActorRegistry::new(120);
        let token = registry.register(RegisterRequest {
            actor_type: "developer".into(),
            hostname: "testhost".into(),
            pid: 9999,
            project: "ontoref".into(),
            ..Default::default()
        });

        assert_eq!(registry.count(), 1);
        assert!(registry.get(&token).is_some());

        assert!(registry.deregister(&token));
        assert_eq!(registry.count(), 0);
    }

    #[test]
    fn list_for_project_filters() {
        let registry = ActorRegistry::new(120);
        registry.register(RegisterRequest {
            actor_type: "developer".into(),
            hostname: "h1".into(),
            pid: 1,
            project: "alpha".into(),
            ..Default::default()
        });
        registry.register(RegisterRequest {
            actor_type: "agent".into(),
            hostname: "h2".into(),
            pid: 2,
            project: "beta".into(),
            ..Default::default()
        });

        let alpha_actors = registry.list_for_project("alpha");
        assert_eq!(alpha_actors.len(), 1);
        assert_eq!(alpha_actors[0].1.actor_type, "developer");
    }

    #[test]
    fn pending_notifications_lifecycle() {
        let registry = ActorRegistry::new(120);
        let token = registry.register(RegisterRequest {
            actor_type: "agent".into(),
            hostname: "h".into(),
            pid: 42,
            project: "test".into(),
            ..Default::default()
        });

        registry.increment_pending(&token);
        registry.increment_pending(&token);
        let view = registry.get(&token).expect("session exists");
        assert_eq!(view.pending_notifications, 2);

        registry.clear_pending(&token);
        let view = registry.get(&token).expect("session exists");
        assert_eq!(view.pending_notifications, 0);
    }

    #[test]
    fn get_and_touch_unknown_token_returns_none() {
        let registry = ActorRegistry::new(120);
        assert!(registry.get_and_touch("no-such-token").is_none());
    }

    #[test]
    fn get_and_touch_returns_view_for_registered_token() {
        let registry = ActorRegistry::new(120);
        let token = registry.register(RegisterRequest {
            actor_type: "developer".into(),
            hostname: "host".into(),
            pid: 1,
            project: "proj".into(),
            ..Default::default()
        });
        let view = registry.get_and_touch(&token);
        assert!(view.is_some());
        assert_eq!(view.unwrap().actor_type, "developer");
    }

    #[test]
    fn get_and_touch_updates_last_seen_atomically() {
        let registry = ActorRegistry::new(120);
        let token = registry.register(RegisterRequest {
            actor_type: "agent".into(),
            hostname: "host".into(),
            pid: 2,
            project: "proj".into(),
            ..Default::default()
        });

        // Force a known-old last_seen so we can detect the update.
        registry.sessions.get_mut(&token).unwrap().last_seen = 1_000_000;

        let view = registry.get_and_touch(&token).unwrap();

        // The returned view must already reflect the updated last_seen.
        assert!(
            view.last_seen > 1_000_000,
            "returned view must carry updated last_seen, got {}",
            view.last_seen
        );
        // The in-registry value must match what was returned — no drift between
        // the write inside get_mut and the subsequent ActorSessionView snapshot.
        let stored = registry.get(&token).unwrap();
        assert_eq!(stored.last_seen, view.last_seen);
    }

    #[cfg(unix)]
    #[test]
    fn sweep_reaps_dead_processes() {
        let registry = ActorRegistry::new(120);
        let local_host = registry.local_hostname.clone();

        // PID 999999 almost certainly does not exist
        registry.register(RegisterRequest {
            actor_type: "agent".into(),
            hostname: local_host,
            pid: 999_999,
            project: "test".into(),
            ..Default::default()
        });

        let reaped = registry.sweep_stale();
        assert_eq!(reaped.len(), 1);
        assert_eq!(registry.count(), 0);
    }
}