deepslate 0.2.0

//! Named server registry with ordered try-list and forced-host routing for
//! initial server selection.

use std::collections::HashMap;
use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

use arc_swap::ArcSwap;
use tracing::info;

/// A compile-time server identifier for use with the builder API.
///
/// `ServerId` uses `&'static str` fields, making it `const`-constructible and
/// `Copy`. Declare your servers as constants and pass them to
/// [`ProxyBuilder::server`] and [`ProxyBuilder::try_servers`] for type-safe
/// server references that prevent accidental typos.
///
/// ```rust
/// use deepslate::ServerId;
///
/// const LOBBY: ServerId = ServerId::new("lobby", "127.0.0.1:25566");
/// const SURVIVAL: ServerId = ServerId::new("survival", "127.0.0.1:25567");
/// ```
#[derive(Debug, Clone, Copy)]
pub struct ServerId {
    /// Unique identifier (e.g., "lobby", "survival").
    pub id: &'static str,
    /// Upstream address (host:port).
    pub addr: &'static str,
}

impl ServerId {
    /// Create a new server identifier.
    #[must_use]
    pub const fn new(id: &'static str, addr: &'static str) -> Self {
        Self { id, addr }
    }
}

/// A single upstream backend server.
#[derive(Debug, Clone)]
pub struct Server {
    /// Unique identifier (e.g., "lobby", "survival").
    pub id: String,
    /// Upstream address (host:port).
    pub addr: String,
}

impl Server {
    /// Create a new server with the given ID and address.
    #[must_use]
    pub fn new(id: impl Into<String>, addr: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            addr: addr.into(),
        }
    }
}

impl From<&ServerId> for Server {
    fn from(server_id: &ServerId) -> Self {
        Self {
            id: server_id.id.to_owned(),
            addr: server_id.addr.to_owned(),
        }
    }
}

/// Thread-safe registry of named backend servers.
///
/// Servers are registered by ID and looked up by name. A configurable
/// "try" list determines the order in which servers are tried for a
/// player's initial connection. An optional forced-hosts map overrides
/// the try list on a per-hostname basis.
///
/// Uses [`ArcSwap`] internally for wait-free reads, making lookups during
/// connection setup contention-free even under high connection rates.
pub struct ServerRegistry {
    servers: ArcSwap<Vec<Server>>,
    try_order: ArcSwap<Vec<String>>,
    forced_hosts: ArcSwap<HashMap<String, Vec<String>>>,
}

impl ServerRegistry {
    /// Create an empty server registry.
    #[must_use]
    pub fn new() -> Self {
        Self {
            servers: ArcSwap::from_pointee(Vec::new()),
            try_order: ArcSwap::from_pointee(Vec::new()),
            forced_hosts: ArcSwap::from_pointee(HashMap::new()),
        }
    }

    /// Register a new server.
    /// Returns `true` if the server was added, `false` if a server with the
    /// same ID already exists.
    pub fn register(&self, server: &Server) -> bool {
        let added = AtomicBool::new(false);
        self.servers.rcu(|current| {
            if current.iter().any(|s| s.id == server.id) {
                Vec::clone(current)
            } else {
                added.store(true, Ordering::Relaxed);
                let mut new = Vec::clone(current);
                new.push(server.clone());
                new
            }
        });

        let was_added = added.load(Ordering::Relaxed);
        if was_added {
            info!(id = server.id, addr = server.addr, "registered server");
        }
        was_added
    }

    /// Deregister a server by ID.
    /// Returns the removed server, or `None` if not found.
    pub fn deregister(&self, id: &str) -> Option<Server> {
        let current = self.servers.load();
        let pos = current.iter().position(|s| s.id == id)?;
        let removed = current[pos].clone();
        self.servers.rcu(|current| {
            let mut new = Vec::clone(current);
            if let Some(pos) = new.iter().position(|s| s.id == id) {
                new.remove(pos);
            }
            new
        });
        Some(removed)
    }

    /// Look up a server by ID.
    pub fn get(&self, id: &str) -> Option<Server> {
        let servers = self.servers.load();
        servers.iter().find(|s| s.id == id).cloned()
    }

    /// List all registered servers.
    pub fn list(&self) -> Vec<Server> {
        self.servers.load().as_ref().clone()
    }

    /// Set the try order for initial server selection.
    ///
    /// The try list is an ordered list of server IDs. When a player first
    /// connects, the proxy tries each server in order until one is found
    /// in the registry.
    pub fn set_try_order(&self, ids: Vec<String>) {
        self.try_order.store(Arc::new(ids));
    }

    /// Get the current try order.
    pub fn try_order(&self) -> Vec<String> {
        self.try_order.load().as_ref().clone()
    }

    /// Select the first available server from the try list.
    ///
    /// Iterates the try order and returns the first server that exists
    /// in the registry. Returns `None` if no try-list server is registered.
    pub fn select_initial(&self) -> Option<Server> {
        let try_order = self.try_order.load();
        let servers = self.servers.load();

        for id in try_order.iter() {
            if let Some(server) = servers.iter().find(|s| s.id == *id) {
                return Some(server.clone());
            }
        }

        // Fallback: if the try list is empty or has no matches, return the first
        // registered server (if any).
        servers.first().cloned()
    }

    /// Set the forced-hosts map.
    ///
    /// Each key is a hostname (lowercased) and each value is an ordered list
    /// of server IDs to try when a player connects using that hostname.
    pub fn set_forced_hosts(&self, map: HashMap<String, Vec<String>>) {
        self.forced_hosts.store(Arc::new(map));
    }

    /// Get a snapshot of the current forced-hosts map.
    pub fn forced_hosts(&self) -> HashMap<String, Vec<String>> {
        self.forced_hosts.load().as_ref().clone()
    }

    /// Select a server for the given virtual hostname.
    ///
    /// Resolution order:
    /// 1. If `hostname` matches a forced-host entry, try those servers in order.
    /// 2. Otherwise fall back to [`select_initial`](Self::select_initial).
    pub fn select_for_host(&self, hostname: &str) -> Option<Server> {
        let forced = self.forced_hosts.load();
        if let Some(ids) = forced.get(hostname) {
            let servers = self.servers.load();
            for id in ids {
                if let Some(server) = servers.iter().find(|s| s.id == *id) {
                    return Some(server.clone());
                }
            }
            // All forced-host servers are unregistered — fall through to the
            // global try list rather than returning None, so the player still
            // has a chance to connect somewhere.
        }
        self.select_initial()
    }
}

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

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

    #[test]
    fn test_register_and_list() {
        let reg = ServerRegistry::new();
        let server = Server::new("lobby", "127.0.0.1:25565");

        assert!(reg.register(&server));
        assert!(!reg.register(&server)); // Duplicate

        let servers = reg.list();
        assert_eq!(servers.len(), 1);
        assert_eq!(servers[0].id, "lobby");
    }

    #[test]
    fn test_deregister() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));

        let removed = reg.deregister("lobby");
        assert!(removed.is_some());
        assert_eq!(removed.unwrap().id, "lobby");
        assert!(reg.list().is_empty());

        assert!(reg.deregister("nonexistent").is_none());
    }

    #[test]
    fn test_get() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));
        reg.register(&Server::new("survival", "127.0.0.1:25566"));

        let server = reg.get("survival").unwrap();
        assert_eq!(server.addr, "127.0.0.1:25566");

        assert!(reg.get("nonexistent").is_none());
    }

    #[test]
    fn test_try_order_selection() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("survival", "127.0.0.1:25566"));
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));

        // Without a try order, fallback returns the first registered server
        assert_eq!(reg.select_initial().unwrap().id, "survival");

        // With a try order, it respects the order
        reg.set_try_order(vec!["lobby".to_string(), "survival".to_string()]);
        assert_eq!(reg.select_initial().unwrap().id, "lobby");
    }

    #[test]
    fn test_try_order_skips_missing() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("survival", "127.0.0.1:25566"));

        // "lobby" isn't registered, so it skips to "survival"
        reg.set_try_order(vec!["lobby".to_string(), "survival".to_string()]);
        assert_eq!(reg.select_initial().unwrap().id, "survival");
    }

    #[test]
    fn test_select_initial_empty() {
        let reg = ServerRegistry::new();
        assert!(reg.select_initial().is_none());
    }

    #[test]
    fn test_try_order_roundtrip() {
        let reg = ServerRegistry::new();
        let order = vec!["a".to_string(), "b".to_string(), "c".to_string()];
        reg.set_try_order(order.clone());
        assert_eq!(reg.try_order(), order);
    }

    #[test]
    fn test_forced_hosts_roundtrip() {
        let reg = ServerRegistry::new();
        let mut map = HashMap::new();
        map.insert("pvp.example.com".to_string(), vec!["pvp".to_string()]);
        reg.set_forced_hosts(map.clone());
        assert_eq!(reg.forced_hosts(), map);
    }

    #[test]
    fn test_select_for_host_exact_match() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));
        reg.register(&Server::new("pvp", "127.0.0.1:25566"));
        reg.set_try_order(vec!["lobby".to_string()]);

        let mut forced = HashMap::new();
        forced.insert("pvp.example.com".to_string(), vec!["pvp".to_string()]);
        reg.set_forced_hosts(forced);

        assert_eq!(reg.select_for_host("pvp.example.com").unwrap().id, "pvp");
    }

    #[test]
    fn test_select_for_host_falls_back_to_try_order() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));
        reg.register(&Server::new("pvp", "127.0.0.1:25566"));
        reg.set_try_order(vec!["lobby".to_string()]);

        let mut forced = HashMap::new();
        forced.insert("pvp.example.com".to_string(), vec!["pvp".to_string()]);
        reg.set_forced_hosts(forced);

        // Unknown hostname falls back to try_order
        assert_eq!(
            reg.select_for_host("unknown.example.com").unwrap().id,
            "lobby"
        );
    }

    #[test]
    fn test_select_for_host_skips_missing_servers() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));
        reg.register(&Server::new("pvp2", "127.0.0.1:25567"));
        reg.set_try_order(vec!["lobby".to_string()]);

        let mut forced = HashMap::new();
        forced.insert(
            "pvp.example.com".to_string(),
            vec!["pvp-gone".to_string(), "pvp2".to_string()],
        );
        reg.set_forced_hosts(forced);

        // "pvp-gone" isn't registered, so it skips to "pvp2"
        assert_eq!(reg.select_for_host("pvp.example.com").unwrap().id, "pvp2");
    }

    #[test]
    fn test_select_for_host_all_forced_missing_falls_back() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));
        reg.set_try_order(vec!["lobby".to_string()]);

        let mut forced = HashMap::new();
        forced.insert(
            "pvp.example.com".to_string(),
            vec!["gone1".to_string(), "gone2".to_string()],
        );
        reg.set_forced_hosts(forced);

        // All forced servers missing — fall back to try order
        assert_eq!(reg.select_for_host("pvp.example.com").unwrap().id, "lobby");
    }

    #[test]
    fn test_select_for_host_empty_hostname() {
        let reg = ServerRegistry::new();
        reg.register(&Server::new("lobby", "127.0.0.1:25565"));
        reg.set_try_order(vec!["lobby".to_string()]);

        // Empty hostname — no forced host match, falls back to try order
        assert_eq!(reg.select_for_host("").unwrap().id, "lobby");
    }
}