latticeon 0.1.0

//! Universe: top-level ECS container.
//!
//! Holds all archetypes, entity allocator, component registry, and entity -> (archetype, row) mapping.

use std::collections::HashMap;

use crate::ecs::archetype::{Archetype, ArchetypeBuilder, ArchetypeId, ComponentBundle};
use crate::ecs::component::{Component, ComponentId, ComponentIdRegistry};
use crate::ecs::entity::{Entity, EntityAllocator};
use crate::ecs::storage::DEFAULT_BLOCK_SIZE;

// ---------------------------------------------------------------------------
// ArchetypeHandle -- lightweight reference to an archetype within a Universe
// ---------------------------------------------------------------------------

/// Opaque handle to a registered archetype. Cheap to copy and reuse.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ArchetypeHandle {
    index: usize,
}

pub struct Universe {
    block_size: usize,
    registry: ComponentIdRegistry,
    entities: EntityAllocator,
    archetypes: Vec<Archetype>,
    archetype_index: HashMap<ArchetypeId, usize>,
    entity_index: HashMap<Entity, (usize, usize)>,
}

impl Universe {
    pub fn new() -> Self {
        Self::with_block_size(DEFAULT_BLOCK_SIZE)
    }

    pub fn with_block_size(block_size: usize) -> Self {
        Self {
            block_size,
            registry: ComponentIdRegistry::new(),
            entities: EntityAllocator::new(),
            archetypes: Vec::new(),
            archetype_index: HashMap::new(),
            entity_index: HashMap::new(),
        }
    }

    #[inline]
    pub fn block_size(&self) -> usize { self.block_size }

    #[inline]
    pub fn registry(&self) -> &ComponentIdRegistry { &self.registry }

    fn get_or_create_archetype_idx<B: ComponentBundle>(&mut self) -> usize {
        let id = ArchetypeId::from_bundle::<B>(&self.registry);
        if let Some(&idx) = self.archetype_index.get(&id) {
            return idx;
        }
        let metas: Vec<_> = B::column_metas(&self.registry)
            .into_iter()
            .map(|cm| (cm.component_id, cm.meta))
            .collect();
        let arch = Archetype::new(id.clone(), &metas, self.block_size);
        let idx = self.archetypes.len();
        self.archetype_index.insert(id, idx);
        self.archetypes.push(arch);
        idx
    }

    pub fn spawn<B: ComponentBundle>(&mut self, bundle: B) -> Entity {
        let entity = self.entities.alloc();
        let arch_idx = self.get_or_create_archetype_idx::<B>();
        let arch = &mut self.archetypes[arch_idx];
        let row = arch.add_entity_with_bundle(entity, bundle, &self.registry);
        self.entity_index.insert(entity, (arch_idx, row));
        entity
    }

    pub fn spawn_empty(&mut self) -> Entity {
        self.spawn(())
    }

    pub fn despawn(&mut self, entity: Entity) -> bool {
        let (arch_idx, row) = match self.entity_index.remove(&entity) {
            Some(x) => x,
            None => return false,
        };
        if !self.entities.is_valid(entity) {
            return false;
        }
        self.entities.free(entity);
        let arch = &mut self.archetypes[arch_idx];
        let last_index = arch.len() - 1;
        let moved_entity = (row != last_index)
            .then(|| arch.entity_at(last_index))
            .flatten();
        arch.remove_entity(row);
        if let Some(moved) = moved_entity {
            self.entity_index.insert(moved, (arch_idx, row));
        }
        true
    }

    pub fn get<T: Component>(&self, entity: Entity) -> Option<&T> {
        let &(arch_idx, row) = self.entity_index.get(&entity)?;
        let arch = self.archetypes.get(arch_idx)?;
        arch.get_comp::<T>(row, &self.registry)
    }

    pub fn get_mut<T: Component>(&mut self, entity: Entity) -> Option<&mut T> {
        let &(arch_idx, row) = self.entity_index.get(&entity)?;
        let arch = self.archetypes.get_mut(arch_idx)?;
        arch.get_comp_mut::<T>(row, &self.registry)
    }

    pub fn contains_entity(&self, entity: Entity) -> bool {
        self.entity_index.contains_key(&entity) && self.entities.is_valid(entity)
    }

    pub fn entity_count(&self) -> usize {
        self.entity_index.len()
    }

    pub fn archetypes(&self) -> impl Iterator<Item = &Archetype> {
        self.archetypes.iter()
    }

    pub fn archetypes_mut(&mut self) -> impl Iterator<Item = &mut Archetype> {
        self.archetypes.iter_mut()
    }

    pub fn get_archetype(&self, id: &ArchetypeId) -> Option<&Archetype> {
        let &idx = self.archetype_index.get(id)?;
        self.archetypes.get(idx)
    }

    pub fn get_archetype_mut(&mut self, id: &ArchetypeId) -> Option<&mut Archetype> {
        let &idx = self.archetype_index.get(id)?;
        self.archetypes.get_mut(idx)
    }

    // -- ArchetypeHandle API --

    /// Ensure the archetype for bundle `B` exists and return a handle to it.
    pub fn register_archetype<B: ComponentBundle>(&mut self) -> ArchetypeHandle {
        let idx = self.get_or_create_archetype_idx::<B>();
        ArchetypeHandle { index: idx }
    }

    /// Register an archetype built dynamically via `ArchetypeBuilder`.
    /// If an archetype with the same component set already exists, returns its handle.
    pub fn register_archetype_dynamic(&mut self, builder: ArchetypeBuilder) -> ArchetypeHandle {
        let id = builder.archetype_id();
        if let Some(&idx) = self.archetype_index.get(&id) {
            return ArchetypeHandle { index: idx };
        }
        let arch = builder.build(self.block_size);
        let idx = self.archetypes.len();
        self.archetype_index.insert(id, idx);
        self.archetypes.push(arch);
        ArchetypeHandle { index: idx }
    }

    /// Begin spawning an entity into the archetype referenced by `handle`.
    ///
    /// Returns an `EntityBuilder` that lets you write component values one by one,
    /// then finalize with `.finish()`.
    pub fn spawn_at(&mut self, handle: ArchetypeHandle) -> EntityBuilder<'_> {
        let entity = self.entities.alloc();
        let arch = &mut self.archetypes[handle.index];
        let row = arch.push_entity(entity);
        self.entity_index.insert(entity, (handle.index, row));
        EntityBuilder {
            universe: self,
            arch_idx: handle.index,
            entity,
            row,
        }
    }

    /// Look up the handle for an already-registered archetype id.
    pub fn handle_for(&self, id: &ArchetypeId) -> Option<ArchetypeHandle> {
        self.archetype_index.get(id).map(|&index| ArchetypeHandle { index })
    }
}

// ---------------------------------------------------------------------------
// EntityBuilder
// ---------------------------------------------------------------------------

/// Builder returned by `Universe::spawn_at` for writing components one by one.
pub struct EntityBuilder<'u> {
    universe: &'u mut Universe,
    arch_idx: usize,
    entity: Entity,
    row: usize,
}

impl<'u> EntityBuilder<'u> {
    /// Write a component value. The component type must be part of the archetype's layout.
    ///
    /// # Panics
    /// Panics if `T` is not present in the archetype's column layout.
    pub fn set<T: Component>(self, value: T) -> Self {
        let cid = self.universe.registry.id_for::<T>();
        let arch = &mut self.universe.archetypes[self.arch_idx];
        unsafe { arch.write_component(self.row, cid, value); }
        self
    }

    /// Write a component value by its runtime `ComponentId`.
    ///
    /// # Safety
    /// `T` must match the type registered under `component_id`.
    pub unsafe fn set_by_id<T: Component>(self, component_id: ComponentId, value: T) -> Self {
        let arch = &mut self.universe.archetypes[self.arch_idx];
        arch.write_component(self.row, component_id, value);
        self
    }

    /// Finalize the entity and return its handle.
    pub fn finish(self) -> Entity {
        self.entity
    }
}

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

unsafe impl Send for Universe {}
unsafe impl Sync for Universe {}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ecs::component::Component;

    #[derive(Debug, Clone, PartialEq)]
    struct Pos(i32, i32);
    impl Component for Pos {}

    #[derive(Debug, Clone, PartialEq)]
    struct Vel(i32, i32);
    impl Component for Vel {}

    #[derive(Debug, Clone, PartialEq)]
    struct Hp(i32);
    impl Component for Hp {}

    #[test]
    fn spawn_and_get() {
        let mut u = Universe::new();
        let e = u.spawn((Pos(1, 2), Vel(3, 4)));
        assert_eq!(u.get::<Pos>(e), Some(&Pos(1, 2)));
        assert_eq!(u.get::<Vel>(e), Some(&Vel(3, 4)));
        u.despawn(e);
        assert!(u.get::<Pos>(e).is_none());
    }

    #[test]
    fn register_archetype_static() {
        let mut u = Universe::new();
        let h1 = u.register_archetype::<(Pos, Vel)>();
        let h2 = u.register_archetype::<(Pos, Vel)>();
        assert_eq!(h1, h2);
    }

    #[test]
    fn spawn_at_handle() {
        let mut u = Universe::new();
        let handle = u.register_archetype::<(Pos, Vel, Hp)>();

        let e = u.spawn_at(handle)
            .set(Pos(1, 2))
            .set(Vel(3, 4))
            .set(Hp(100))
            .finish();

        assert_eq!(u.get::<Pos>(e), Some(&Pos(1, 2)));
        assert_eq!(u.get::<Vel>(e), Some(&Vel(3, 4)));
        assert_eq!(u.get::<Hp>(e), Some(&Hp(100)));
        assert_eq!(u.entity_count(), 1);
    }

    #[test]
    fn spawn_at_multiple() {
        let mut u = Universe::new();
        let handle = u.register_archetype::<(Pos, Vel)>();

        let e1 = u.spawn_at(handle).set(Pos(1, 1)).set(Vel(0, 0)).finish();
        let e2 = u.spawn_at(handle).set(Pos(2, 2)).set(Vel(1, 1)).finish();

        assert_eq!(u.entity_count(), 2);
        assert_eq!(u.get::<Pos>(e1), Some(&Pos(1, 1)));
        assert_eq!(u.get::<Pos>(e2), Some(&Pos(2, 2)));
    }

    #[test]
    fn register_archetype_dynamic_builder() {
        let mut u = Universe::new();

        let mut builder = ArchetypeBuilder::new();
        builder.add_component::<Pos>(u.registry()).add_component::<Hp>(u.registry());
        let handle = u.register_archetype_dynamic(builder);

        let e = u.spawn_at(handle)
            .set(Pos(10, 20))
            .set(Hp(50))
            .finish();

        assert_eq!(u.get::<Pos>(e), Some(&Pos(10, 20)));
        assert_eq!(u.get::<Hp>(e), Some(&Hp(50)));
    }

    #[test]
    fn dynamic_and_static_same_archetype() {
        let mut u = Universe::new();

        let h_static = u.register_archetype::<(Pos, Vel)>();
        let mut builder = ArchetypeBuilder::new();
        builder.add_component::<Pos>(u.registry()).add_component::<Vel>(u.registry());
        let h_dynamic = u.register_archetype_dynamic(builder);

        assert_eq!(h_static, h_dynamic);
    }

    #[test]
    fn despawn_after_spawn_at() {
        let mut u = Universe::new();
        let handle = u.register_archetype::<(Pos, Vel)>();

        let e = u.spawn_at(handle).set(Pos(1, 2)).set(Vel(3, 4)).finish();
        assert!(u.contains_entity(e));
        u.despawn(e);
        assert!(!u.contains_entity(e));
    }
}