settled-reach/server/src/knowledge/registry.rs

//! Entity ID infrastructure (Q-019, D-041).
//!
//! StableEntityId component and EntityRegistry resource for
//! bidirectional StableId <-> Entity mapping.

use bevy_ecs::prelude::*;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

use super::types::StableId;

/// Stable entity ID component. Assigned once at spawn, never changes.
/// Serialized with entity for save/load.
#[derive(Component, Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct StableEntityId(pub StableId);

/// Bidirectional mapping between StableId and bevy Entity.
/// Resource, updated on spawn/despawn.
#[derive(Resource, Debug)]
pub struct EntityRegistry {
    by_stable_id: BTreeMap<StableId, Entity>,
    by_entity: BTreeMap<Entity, StableId>,
    next_id: u64,
}

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

impl EntityRegistry {
    pub fn new(seed_offset: u64) -> Self {
        Self {
            by_stable_id: BTreeMap::new(),
            by_entity: BTreeMap::new(),
            next_id: seed_offset,
        }
    }

    /// Register a new entity and assign a StableId.
    /// Returns the existing StableId if already registered.
    pub fn register(&mut self, entity: Entity) -> StableId {
        if let Some(id) = self.by_entity.get(&entity) {
            return *id;
        }
        let id = StableId(self.next_id);
        self.next_id += 1;
        self.by_stable_id.insert(id, entity);
        self.by_entity.insert(entity, id);
        id
    }

    /// Lookup: StableId -> Entity (for ECS queries).
    pub fn to_entity(&self, id: &StableId) -> Option<Entity> {
        self.by_stable_id.get(id).copied()
    }

    /// Lookup: Entity -> StableId (for knowledge graph keys).
    pub fn to_stable(&self, entity: Entity) -> Option<StableId> {
        self.by_entity.get(&entity).copied()
    }

    /// Remove a despawned entity from the registry.
    /// Called AFTER knowledge cleanup (tombstone pattern).
    pub fn unregister(&mut self, entity: Entity) {
        if let Some(id) = self.by_entity.remove(&entity) {
            self.by_stable_id.remove(&id);
        }
    }

    /// Advance the next StableId counter to `target`.
    /// Used to reserve StableId ranges for entities not yet spawned
    /// (e.g., Gauntlet rooms built in later sprints).
    /// Panics if `target` is less than the current next_id.
    pub fn reserve_up_to(&mut self, target: u64) {
        assert!(
            target >= self.next_id,
            "cannot reserve backwards: next_id={}, target={}",
            self.next_id,
            target
        );
        self.next_id = target;
    }

    /// Register an entity with a specific pre-existing StableId (used during save/load).
    ///
    /// Unlike `register`, this does NOT advance `next_id`. After bulk-registering
    /// all restored entities, call `advance_past(max_stable_id)` so future `register()`
    /// calls produce IDs that don't conflict with the restored set.
    ///
    /// No-op if the entity is already mapped to the same `stable_id`.
    /// Panics in debug builds if `stable_id` is already mapped to a different entity.
    pub fn register_existing(&mut self, entity: Entity, stable_id: StableId) {
        if let Some(&existing) = self.by_stable_id.get(&stable_id) {
            debug_assert_eq!(
                existing, entity,
                "register_existing: StableId {:?} already mapped to a different entity",
                stable_id
            );
            return;
        }
        self.by_stable_id.insert(stable_id, entity);
        self.by_entity.insert(entity, stable_id);
    }

    /// Advance `next_id` past `id` so future `register()` calls don't conflict.
    ///
    /// Unlike `reserve_up_to`, this never panics: if the counter is already past `id`,
    /// this is a no-op. Use after `register_existing` bulk-load to position the counter.
    pub fn advance_past(&mut self, id: u64) {
        let target = id.saturating_add(1);
        if target > self.next_id {
            self.next_id = target;
        }
    }

    /// Number of registered entities.
    pub fn len(&self) -> usize {
        self.by_entity.len()
    }

    /// Whether the registry is empty.
    pub fn is_empty(&self) -> bool {
        self.by_entity.is_empty()
    }
}

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

    #[test]
    fn register_assigns_sequential_ids() {
        let mut world = World::new();
        let e1 = world.spawn_empty().id();
        let e2 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(0);
        let id1 = registry.register(e1);
        let id2 = registry.register(e2);

        assert_eq!(id1, StableId(0));
        assert_eq!(id2, StableId(1));
    }

    #[test]
    fn register_idempotent() {
        let mut world = World::new();
        let e1 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(0);
        let id_first = registry.register(e1);
        let id_second = registry.register(e1);

        assert_eq!(id_first, id_second);
        assert_eq!(registry.len(), 1);
    }

    #[test]
    fn bidirectional_lookup() {
        let mut world = World::new();
        let e1 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(0);
        let id = registry.register(e1);

        assert_eq!(registry.to_entity(&id), Some(e1));
        assert_eq!(registry.to_stable(e1), Some(id));
    }

    #[test]
    fn unregister_removes_both_directions() {
        let mut world = World::new();
        let e1 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(0);
        let id = registry.register(e1);
        registry.unregister(e1);

        assert_eq!(registry.to_entity(&id), None);
        assert_eq!(registry.to_stable(e1), None);
        assert!(registry.is_empty());
    }

    #[test]
    fn seed_offset_starts_ids_higher() {
        let mut world = World::new();
        let e1 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(1000);
        let id = registry.register(e1);

        assert_eq!(id, StableId(1000));
    }

    #[test]
    fn lookup_unknown_entity_returns_none() {
        let mut world = World::new();
        let e1 = world.spawn_empty().id();

        let registry = EntityRegistry::new(0);
        assert_eq!(registry.to_entity(&StableId(999)), None);
        assert_eq!(registry.to_stable(e1), None);
    }

    // === EntityRegistry lifecycle edge cases (#469) ===

    #[test]
    fn stale_mapping_after_despawn() {
        // #469: Registry returns stale Entity after world despawn.
        // This documents the expected behavior — caller must unregister after despawn.
        let mut world = World::new();
        let e1 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(0);
        let id = registry.register(e1);

        // Despawn from world — registry doesn't know
        world.despawn(e1);

        // Registry still maps the StableId to the (now stale) Entity
        let stale_entity = registry.to_entity(&id);
        assert!(
            stale_entity.is_some(),
            "Registry still holds mapping after world despawn"
        );

        // But the world no longer recognizes the entity
        assert!(
            world.get_entity(stale_entity.unwrap()).is_err(),
            "World rejects stale entity — caller must call unregister()"
        );

        // After proper cleanup, mapping is gone
        registry.unregister(e1);
        assert_eq!(
            registry.to_entity(&id),
            None,
            "Mapping gone after unregister"
        );
    }

    #[test]
    fn register_after_unregister_assigns_new_id() {
        // #469: Re-registering the same entity after unregister gets a new StableId.
        // StableId counter is monotonic — never recycles.
        let mut world = World::new();
        let e1 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(0);
        let id_first = registry.register(e1);
        assert_eq!(id_first, StableId(0));

        registry.unregister(e1);

        let id_second = registry.register(e1);
        assert_ne!(
            id_first, id_second,
            "Re-registration must assign a new StableId"
        );
        assert_eq!(id_second, StableId(1), "Counter advances monotonically");
        assert_eq!(registry.len(), 1);

        // New mapping is bidirectionally correct
        assert_eq!(registry.to_entity(&id_second), Some(e1));
        assert_eq!(registry.to_stable(e1), Some(id_second));

        // Old StableId no longer resolves
        assert_eq!(
            registry.to_entity(&id_first),
            None,
            "Old StableId must not resolve"
        );
    }

    #[test]
    fn reserve_up_to_advances_counter() {
        let mut world = World::new();
        let mut registry = EntityRegistry::new(0);

        let e1 = world.spawn_empty().id();
        let id1 = registry.register(e1);
        assert_eq!(id1, StableId(0));

        // Reserve through 5 (skip IDs 1-4)
        registry.reserve_up_to(5);

        let e2 = world.spawn_empty().id();
        let id2 = registry.register(e2);
        assert_eq!(id2, StableId(5), "next ID after reserve should be 5");
    }

    #[test]
    #[should_panic(expected = "cannot reserve backwards")]
    fn reserve_up_to_panics_on_backwards() {
        let mut registry = EntityRegistry::new(10);
        registry.reserve_up_to(5);
    }

    #[test]
    fn unregister_unknown_entity_is_noop() {
        // #469: Unregistering an entity that was never registered must not panic.
        let mut world = World::new();
        let e1 = world.spawn_empty().id();
        let e2 = world.spawn_empty().id();

        let mut registry = EntityRegistry::new(0);
        registry.register(e1);

        // Unregister e2 which was never registered — should be a no-op
        registry.unregister(e2);

        // e1's registration is unaffected
        assert_eq!(registry.len(), 1);
        assert_eq!(registry.to_stable(e1), Some(StableId(0)));
    }
}