//! Archetype: set of entities with the same component layout.
//!
//! Storage uses SoA-in-block: each block is one contiguous `std::alloc::alloc` allocation
//! with columns laid out sequentially: `[Entity..] [CompA..] [CompB..]`.

use std::alloc::{self, Layout};
use std::hash::{Hash, Hasher};
use std::ops::Range;
use std::ptr::{self, NonNull};

use smallvec::SmallVec;

use crate::ecs::component::{Component, ComponentId, ComponentIdRegistry, ComponentMeta};
use crate::ecs::entity::Entity;

// ---------------------------------------------------------------------------
// ArchetypeId (dynamic bitvector backed by SmallVec)
// ---------------------------------------------------------------------------

/// Inline capacity: 4 words = 256 bits. Spills to heap beyond that.
const INLINE_WORDS: usize = 4;

#[derive(Clone, PartialEq, Eq)]
pub struct ArchetypeId {
    bits: SmallVec<[u64; INLINE_WORDS]>,
}

impl Hash for ArchetypeId {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        let significant = self.significant_len();
        self.bits[..significant].hash(state);
    }
}

impl Default for ArchetypeId {
    fn default() -> Self { Self::empty() }
}

impl std::fmt::Debug for ArchetypeId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ArchetypeId(")?;
        let sig = self.significant_len();
        if sig == 0 {
            write!(f, "0")?;
        } else {
            for (i, &word) in self.bits[..sig].iter().rev().enumerate() {
                if i > 0 { write!(f, "_")?; }
                write!(f, "{:016x}", word)?;
            }
        }
        write!(f, ")")
    }
}

impl ArchetypeId {
    #[inline]
    pub fn empty() -> Self {
        Self { bits: SmallVec::new() }
    }

    /// Number of words excluding trailing zeros, for consistent hashing.
    #[inline]
    fn significant_len(&self) -> usize {
        let mut len = self.bits.len();
        while len > 0 && self.bits[len - 1] == 0 {
            len -= 1;
        }
        len
    }

    #[inline]
    fn ensure_word(&mut self, word_idx: usize) {
        if word_idx >= self.bits.len() {
            self.bits.resize(word_idx + 1, 0);
        }
    }

    #[inline]
    pub fn set(&mut self, id: ComponentId) {
        let idx = id.0 as usize;
        let word_idx = idx / 64;
        self.ensure_word(word_idx);
        self.bits[word_idx] |= 1u64 << (idx % 64);
    }

    #[inline]
    pub fn single(id: ComponentId) -> Self {
        let mut aid = Self::empty();
        aid.set(id);
        aid
    }

    #[inline]
    pub fn contains_all(&self, mask: &Self) -> bool {
        for (i, &m) in mask.bits.iter().enumerate() {
            let s = self.bits.get(i).copied().unwrap_or(0);
            if s & m != m {
                return false;
            }
        }
        true
    }

    #[inline]
    pub fn has(&self, id: ComponentId) -> bool {
        let idx = id.0 as usize;
        let word_idx = idx / 64;
        match self.bits.get(word_idx) {
            Some(&word) => word & (1u64 << (idx % 64)) != 0,
            None => false,
        }
    }

    #[inline]
    pub fn union(&self, other: &Self) -> Self {
        let max_len = self.bits.len().max(other.bits.len());
        let mut result = SmallVec::with_capacity(max_len);
        for i in 0..max_len {
            let a = self.bits.get(i).copied().unwrap_or(0);
            let b = other.bits.get(i).copied().unwrap_or(0);
            result.push(a | b);
        }
        Self { bits: result }
    }

    pub fn from_bundle<B: ComponentBundle>(registry: &ComponentIdRegistry) -> Self {
        B::archetype_id(registry)
    }

    #[inline]
    pub fn count(&self) -> u32 {
        self.bits.iter().map(|w| w.count_ones()).sum()
    }

    pub fn iter_set(&self) -> impl Iterator<Item = ComponentId> + '_ {
        self.bits.iter().enumerate().flat_map(|(word_idx, &word)| {
            let base = (word_idx * 64) as u32;
            (0..64u32).filter_map(move |bit| {
                if word & (1u64 << bit) != 0 {
                    Some(ComponentId(base + bit))
                } else {
                    None
                }
            })
        })
    }
}

// ---------------------------------------------------------------------------
// ColumnSlice -- describes one column's position inside a block
// ---------------------------------------------------------------------------

#[derive(Clone, Copy)]
pub struct ColumnSlice {
    pub component_id: ComponentId,
    /// Byte offset of this column's first element from block start.
    pub offset: usize,
    pub elem_size: usize,
    pub elem_align: usize,
    pub drop_fn: Option<unsafe fn(*mut u8)>,
}

// ---------------------------------------------------------------------------
// BlockLayout -- SoA layout within a single raw block
// ---------------------------------------------------------------------------

pub struct BlockLayout {
    pub entity_offset: usize,
    pub columns: Vec<ColumnSlice>,
    pub capacity: usize,
    pub block_bytes: usize,
    pub block_align: usize,
}

fn align_up(val: usize, align: usize) -> usize {
    (val + align - 1) & !(align - 1)
}

impl BlockLayout {
    /// Compute the SoA layout for a block of `target_bytes` bytes.
    ///
    /// Layout: `[Entity * cap] [ColA * cap] [ColB * cap] ...`
    /// Each column section starts at a properly aligned offset.
    pub fn new(metas: &[(ComponentId, ComponentMeta)], target_bytes: usize) -> Self {
        if metas.is_empty() {
            return Self {
                entity_offset: 0,
                columns: Vec::new(),
                capacity: target_bytes / std::mem::size_of::<Entity>().max(1),
                block_bytes: target_bytes,
                block_align: std::mem::align_of::<Entity>(),
            };
        }

        let entity_size = std::mem::size_of::<Entity>();
        let entity_align = std::mem::align_of::<Entity>();

        // Compute bytes per entity (entity + all components), ignoring alignment padding for capacity estimate.
        let per_entity_min: usize = entity_size + metas.iter().map(|(_, m)| m.size).sum::<usize>();
        let capacity = if per_entity_min == 0 { 1 } else { (target_bytes / per_entity_min).max(1) };

        let mut max_align = entity_align;
        for (_, m) in metas {
            max_align = max_align.max(m.align);
        }

        // Entity column at offset 0
        let entity_offset = 0usize;
        let mut cursor = entity_size * capacity;

        let mut columns = Vec::with_capacity(metas.len());
        for &(cid, ref meta) in metas {
            cursor = align_up(cursor, meta.align);
            columns.push(ColumnSlice {
                component_id: cid,
                offset: cursor,
                elem_size: meta.size,
                elem_align: meta.align,
                drop_fn: meta.drop_fn,
            });
            cursor += meta.size * capacity;
        }

        let block_bytes = align_up(cursor, max_align);

        Self {
            entity_offset,
            columns,
            capacity,
            block_bytes,
            block_align: max_align,
        }
    }

    /// Find the column slice for a given component id.
    #[inline]
    pub fn column_for(&self, id: ComponentId) -> Option<&ColumnSlice> {
        self.columns.iter().find(|c| c.component_id == id)
    }
}

// ---------------------------------------------------------------------------
// RawBlock -- one contiguous allocation
// ---------------------------------------------------------------------------

pub struct RawBlock {
    ptr: NonNull<u8>,
    alloc_layout: Layout,
}

impl RawBlock {
    fn new(layout: &BlockLayout) -> Self {
        if layout.block_bytes == 0 {
            return Self {
                ptr: NonNull::dangling(),
                alloc_layout: Layout::from_size_align(0, 1).unwrap(),
            };
        }
        let alloc_layout = Layout::from_size_align(layout.block_bytes, layout.block_align)
            .expect("invalid block layout");
        let ptr = unsafe { alloc::alloc_zeroed(alloc_layout) };
        let ptr = NonNull::new(ptr).expect("allocation failed");
        Self { ptr, alloc_layout }
    }

    #[inline]
    fn base(&self) -> *mut u8 {
        self.ptr.as_ptr()
    }
}

impl Drop for RawBlock {
    fn drop(&mut self) {
        if self.alloc_layout.size() > 0 {
            unsafe { alloc::dealloc(self.ptr.as_ptr(), self.alloc_layout); }
        }
    }
}

unsafe impl Send for RawBlock {}
unsafe impl Sync for RawBlock {}

// ---------------------------------------------------------------------------
// Archetype
// ---------------------------------------------------------------------------

pub struct Archetype {
    id: ArchetypeId,
    layout: BlockLayout,
    blocks: Vec<RawBlock>,
    len: usize,
    block_size: usize,
}

impl Archetype {
    /// Create from a precomputed BlockLayout.
    pub fn with_layout(id: ArchetypeId, layout: BlockLayout, block_size: usize) -> Self {
        Self { id, layout, blocks: Vec::new(), len: 0, block_size }
    }

    /// Create from column metas and target block size.
    pub fn new(id: ArchetypeId, metas: &[(ComponentId, ComponentMeta)], block_size: usize) -> Self {
        let layout = BlockLayout::new(metas, block_size);
        Self::with_layout(id, layout, block_size)
    }

    #[inline]
    pub fn id(&self) -> &ArchetypeId { &self.id }

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

    #[inline]
    pub fn is_empty(&self) -> bool { self.len == 0 }

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

    #[inline]
    pub fn layout(&self) -> &BlockLayout { &self.layout }

    #[inline]
    pub fn has_component(&self, id: ComponentId) -> bool { self.id.has(id) }

    /// Number of allocated blocks currently in use (containing at least one entity).
    #[inline]
    pub fn block_count(&self) -> usize {
        if self.len == 0 || self.layout.capacity == 0 {
            return 0;
        }
        (self.len + self.layout.capacity - 1) / self.layout.capacity
    }

    /// Entity capacity per block.
    #[inline]
    pub fn block_capacity(&self) -> usize {
        self.layout.capacity
    }

    /// Global row range for the given block index.
    /// The last block may be partially filled.
    #[inline]
    pub fn block_range(&self, block_idx: usize) -> Range<usize> {
        let cap = self.layout.capacity;
        let start = block_idx * cap;
        let end = (start + cap).min(self.len);
        start..end
    }

    /// Raw base pointer for block `bi`.
    ///
    /// # Safety
    /// Caller must ensure `bi < self.block_count()` and must not alias with
    /// other mutable access to the same block concurrently.
    #[inline]
    pub unsafe fn block_base_ptr(&self, bi: usize) -> *mut u8 {
        self.blocks[bi].base()
    }

    // -- index helpers --

    #[inline]
    fn block_and_offset(&self, index: usize) -> (usize, usize) {
        let cap = self.layout.capacity;
        (index / cap, index % cap)
    }

    fn ensure_capacity(&mut self) {
        let cap = self.layout.capacity;
        if cap == 0 { return; }
        let needed_blocks = (self.len / cap) + 1;
        while self.blocks.len() < needed_blocks {
            self.blocks.push(RawBlock::new(&self.layout));
        }
    }

    // -- entity access --

    #[inline]
    pub fn entity_at(&self, index: usize) -> Option<Entity> {
        if index >= self.len { return None; }
        let (bi, off) = self.block_and_offset(index);
        let ptr = unsafe {
            self.blocks[bi].base()
                .add(self.layout.entity_offset)
                .add(off * std::mem::size_of::<Entity>()) as *const Entity
        };
        Some(unsafe { ptr::read(ptr) })
    }

    /// Raw pointer to entity slot (for swap-remove).
    #[inline]
    unsafe fn entity_ptr(&self, bi: usize, off: usize) -> *mut Entity {
        self.blocks[bi].base()
            .add(self.layout.entity_offset)
            .add(off * std::mem::size_of::<Entity>()) as *mut Entity
    }

    // -- component access --

    /// Get a reference to component T at the given global index.
    ///
    /// # Safety
    /// Caller ensures T matches the actual stored type at `component_id`.
    #[inline]
    pub unsafe fn get_component<T: Component>(&self, index: usize, component_id: ComponentId) -> Option<&T> {
        if index >= self.len { return None; }
        let col = self.layout.column_for(component_id)?;
        let (bi, off) = self.block_and_offset(index);
        let ptr = self.blocks[bi].base().add(col.offset).add(off * col.elem_size) as *const T;
        Some(&*ptr)
    }

    /// Get a mutable reference to component T at the given global index.
    ///
    /// # Safety
    /// Caller ensures T matches the actual stored type at `component_id`.
    #[inline]
    pub unsafe fn get_component_mut<T: Component>(&mut self, index: usize, component_id: ComponentId) -> Option<&mut T> {
        if index >= self.len { return None; }
        let col = self.layout.column_for(component_id)?;
        let (bi, off) = self.block_and_offset(index);
        let ptr = self.blocks[bi].base().add(col.offset).add(off * col.elem_size) as *mut T;
        Some(&mut *ptr)
    }

    /// Convenience: get component by registry lookup.
    pub fn get_comp<T: Component>(&self, index: usize, registry: &ComponentIdRegistry) -> Option<&T> {
        let cid = registry.get::<T>()?;
        unsafe { self.get_component(index, cid) }
    }

    /// Convenience: get mutable component by registry lookup.
    pub fn get_comp_mut<T: Component>(&mut self, index: usize, registry: &ComponentIdRegistry) -> Option<&mut T> {
        let cid = registry.get::<T>()?;
        unsafe { self.get_component_mut(index, cid) }
    }

    // -- push --

    /// Reserve the next slot. Writes the entity, returns the global index.
    /// The caller must then write all component data via `write_component`.
    pub fn push_entity(&mut self, entity: Entity) -> usize {
        let index = self.len;
        self.len += 1;
        self.ensure_capacity();
        let (bi, off) = self.block_and_offset(index);
        unsafe {
            ptr::write(self.entity_ptr(bi, off), entity);
        }
        index
    }

    /// Write a component value into the slot at `index`.
    ///
    /// # Safety
    /// - `index` must have been returned by a prior `push_entity` call (not yet populated for this column).
    /// - T must match the type registered under `component_id`.
    pub unsafe fn write_component<T: Component>(&mut self, index: usize, component_id: ComponentId, value: T) {
        let col = self.layout.column_for(component_id).expect("column not in layout");
        let (bi, off) = self.block_and_offset(index);
        let ptr = self.blocks[bi].base().add(col.offset).add(off * col.elem_size) as *mut T;
        ptr::write(ptr, value);
    }

    /// Add an entity with a component bundle in one call.
    pub fn add_entity_with_bundle<B: ComponentBundle>(
        &mut self, entity: Entity, bundle: B, registry: &ComponentIdRegistry,
    ) -> usize {
        let index = self.push_entity(entity);
        bundle.write_components(self, index, registry);
        index
    }

    // -- swap-remove --

    /// Remove entity at `index` by swapping with the last. Returns the removed entity.
    pub fn remove_entity(&mut self, index: usize) -> Option<Entity> {
        if index >= self.len { return None; }
        let last = self.len - 1;

        let (bi, off) = self.block_and_offset(index);
        let removed_entity = unsafe { ptr::read(self.entity_ptr(bi, off)) };

        // Drop components at `index`
        for col in &self.layout.columns {
            if let Some(drop_fn) = col.drop_fn {
                let ptr = unsafe { self.blocks[bi].base().add(col.offset).add(off * col.elem_size) };
                unsafe { drop_fn(ptr); }
            }
        }

        if index != last {
            let (lbi, loff) = self.block_and_offset(last);
            // Move entity
            unsafe {
                let src = self.entity_ptr(lbi, loff);
                let dst = self.entity_ptr(bi, off);
                ptr::copy_nonoverlapping(src as *const u8, dst as *mut u8, std::mem::size_of::<Entity>());
            }
            // Move each column
            for col in &self.layout.columns {
                unsafe {
                    let src = self.blocks[lbi].base().add(col.offset).add(loff * col.elem_size);
                    let dst = self.blocks[bi].base().add(col.offset).add(off * col.elem_size);
                    ptr::copy_nonoverlapping(src, dst, col.elem_size);
                }
            }
        }

        self.len -= 1;
        Some(removed_entity)
    }
}

impl Drop for Archetype {
    fn drop(&mut self) {
        let cap = self.layout.capacity;
        if cap == 0 { return; }

        for block_idx in 0..self.blocks.len() {
            let start = block_idx * cap;
            let end = self.len.min(start + cap);
            if start >= end { break; }

            let base = self.blocks[block_idx].base();
            for col in &self.layout.columns {
                if let Some(drop_fn) = col.drop_fn {
                    for i in 0..(end - start) {
                        unsafe {
                            let ptr = base.add(col.offset).add(i * col.elem_size);
                            drop_fn(ptr);
                        }
                    }
                }
            }
        }
    }
}

unsafe impl Send for Archetype {}
unsafe impl Sync for Archetype {}

// ---------------------------------------------------------------------------
// ComponentBundle
// ---------------------------------------------------------------------------

/// Metadata for one column, provided by a bundle to build the BlockLayout.
pub struct ColumnMeta {
    pub component_id: ComponentId,
    pub meta: ComponentMeta,
}

pub trait ComponentBundle: Send + Sync {
    fn archetype_id(registry: &ComponentIdRegistry) -> ArchetypeId where Self: Sized;
    fn column_metas(registry: &ComponentIdRegistry) -> Vec<ColumnMeta> where Self: Sized;
    fn write_components(self, archetype: &mut Archetype, index: usize, registry: &ComponentIdRegistry);
}

impl ComponentBundle for () {
    fn archetype_id(_registry: &ComponentIdRegistry) -> ArchetypeId { ArchetypeId::empty() }
    fn column_metas(_registry: &ComponentIdRegistry) -> Vec<ColumnMeta> { Vec::new() }
    fn write_components(self, _archetype: &mut Archetype, _index: usize, _registry: &ComponentIdRegistry) {}
}

macro_rules! impl_component_bundle {
    ($($idx:tt => $T:ident),+) => {
        impl<$($T: Component),+> ComponentBundle for ($($T,)+) {
            fn archetype_id(registry: &ComponentIdRegistry) -> ArchetypeId {
                let mut id = ArchetypeId::empty();
                $(id.set(registry.id_for::<$T>());)+
                id
            }
            fn column_metas(registry: &ComponentIdRegistry) -> Vec<ColumnMeta> {
                vec![$(ColumnMeta {
                    component_id: registry.id_for::<$T>(),
                    meta: ComponentMeta::of::<$T>(),
                }),+]
            }
            fn write_components(
                self, archetype: &mut Archetype, index: usize,
                registry: &ComponentIdRegistry,
            ) {
                unsafe {
                    $(archetype.write_component(index, registry.id_for::<$T>(), self.$idx);)+
                }
            }
        }
    };
}

impl_component_bundle!(0 => A);
impl_component_bundle!(0 => A, 1 => B);
impl_component_bundle!(0 => A, 1 => B, 2 => C);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I, 9 => J);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I, 9 => J, 10 => K);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I, 9 => J, 10 => K, 11 => L);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I, 9 => J, 10 => K, 11 => L, 12 => M);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I, 9 => J, 10 => K, 11 => L, 12 => M, 13 => N);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I, 9 => J, 10 => K, 11 => L, 12 => M, 13 => N, 14 => O);
impl_component_bundle!(0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => G, 7 => H, 8 => I, 9 => J, 10 => K, 11 => L, 12 => M, 13 => N, 14 => O, 15 => P);

// ---------------------------------------------------------------------------
// ArchetypeBuilder -- dynamic construction
// ---------------------------------------------------------------------------

/// Accumulates component types at runtime and builds an `Archetype`.
///
/// Unlike the compile-time `ComponentBundle` approach, this allows registering
/// an arbitrary number of component types dynamically.
pub struct ArchetypeBuilder {
    metas: Vec<(ComponentId, ComponentMeta)>,
    archetype_id: ArchetypeId,
}

impl ArchetypeBuilder {
    pub fn new() -> Self {
        Self {
            metas: Vec::new(),
            archetype_id: ArchetypeId::empty(),
        }
    }

    /// Register a single component type. Duplicate types are silently ignored.
    pub fn add_component<T: Component>(&mut self, registry: &ComponentIdRegistry) -> &mut Self {
        let cid = registry.id_for::<T>();
        if !self.archetype_id.has(cid) {
            self.archetype_id.set(cid);
            self.metas.push((cid, ComponentMeta::of::<T>()));
        }
        self
    }

    /// Register all component types from a `ComponentBundle`.
    /// Duplicate types (already present in this builder) are silently ignored.
    pub fn add_components<B: ComponentBundle>(&mut self, registry: &ComponentIdRegistry) -> &mut Self {
        for cm in B::column_metas(registry) {
            if !self.archetype_id.has(cm.component_id) {
                self.archetype_id.set(cm.component_id);
                self.metas.push((cm.component_id, cm.meta));
            }
        }
        self
    }

    pub fn archetype_id(&self) -> ArchetypeId {
        self.archetype_id.clone()
    }

    pub fn metas(&self) -> &[(ComponentId, ComponentMeta)] {
        &self.metas
    }

    pub fn build(self, block_size: usize) -> Archetype {
        Archetype::new(self.archetype_id, &self.metas, block_size)
    }
}

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

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[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 {}

    #[test]
    fn bitvec_contains_all() {
        let mut a = ArchetypeId::empty();
        a.set(ComponentId(0));
        a.set(ComponentId(3));
        a.set(ComponentId(128));

        let mut mask = ArchetypeId::empty();
        mask.set(ComponentId(0));
        mask.set(ComponentId(3));
        assert!(a.contains_all(&mask));

        mask.set(ComponentId(200));
        assert!(!a.contains_all(&mask));
    }

    #[test]
    fn bitvec_identity() {
        let reg = ComponentIdRegistry::new();
        let id1 = <(Pos, Vel)>::archetype_id(&reg);
        let id2 = <(Vel, Pos)>::archetype_id(&reg);
        assert_eq!(id1, id2);
    }

    #[test]
    fn soa_block_add_entity_and_remove() {
        let reg = ComponentIdRegistry::new();
        let aid = <(Pos, Vel)>::archetype_id(&reg);
        let metas: Vec<_> = <(Pos, Vel)>::column_metas(&reg)
            .into_iter()
            .map(|cm| (cm.component_id, cm.meta))
            .collect();
        let mut arch = Archetype::new(aid, &metas, 1024);

        let e0 = Entity::new(0, 0);
        arch.add_entity_with_bundle(e0, (Pos(1, 2), Vel(3, 4)), &reg);
        assert_eq!(arch.len(), 1);
        assert_eq!(arch.entity_at(0), Some(e0));
        assert_eq!(arch.get_comp::<Pos>(0, &reg), Some(&Pos(1, 2)));
        assert_eq!(arch.get_comp::<Vel>(0, &reg), Some(&Vel(3, 4)));

        let e1 = Entity::new(1, 0);
        arch.add_entity_with_bundle(e1, (Pos(5, 6), Vel(7, 8)), &reg);
        assert_eq!(arch.len(), 2);

        let removed = arch.remove_entity(0);
        assert_eq!(removed, Some(e0));
        assert_eq!(arch.len(), 1);
        // e1 was swapped into index 0
        assert_eq!(arch.entity_at(0), Some(e1));
        assert_eq!(arch.get_comp::<Pos>(0, &reg), Some(&Pos(5, 6)));
        assert_eq!(arch.get_comp::<Vel>(0, &reg), Some(&Vel(7, 8)));
    }

    #[test]
    fn soa_block_multi_block() {
        let reg = ComponentIdRegistry::new();
        let aid = <(Pos,)>::archetype_id(&reg);
        let metas: Vec<_> = <(Pos,)>::column_metas(&reg)
            .into_iter()
            .map(|cm| (cm.component_id, cm.meta))
            .collect();
        // Tiny block: Entity(8 bytes) + Pos(8 bytes) = 16 bytes/entity, block_size=64 => cap=4
        let mut arch = Archetype::new(aid, &metas, 64);
        assert_eq!(arch.layout().capacity, 4);

        for i in 0..10u32 {
            let e = Entity::new(i, 0);
            arch.add_entity_with_bundle(e, (Pos(i as i32, 0),), &reg);
        }
        assert_eq!(arch.len(), 10);
        assert!(arch.blocks.len() >= 3); // 10/4 = 3 blocks

        for i in 0..10u32 {
            assert_eq!(arch.entity_at(i as usize), Some(Entity::new(i, 0)));
            assert_eq!(arch.get_comp::<Pos>(i as usize, &reg), Some(&Pos(i as i32, 0)));
        }

        // Remove from middle (cross-block swap)
        arch.remove_entity(2);
        assert_eq!(arch.len(), 9);
        // Entity 9 should now be at index 2
        assert_eq!(arch.entity_at(2), Some(Entity::new(9, 0)));
    }

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

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

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

    #[test]
    fn bundle_5_tuple() {
        let reg = ComponentIdRegistry::new();
        let aid = <(Pos, Vel, Hp, Atk, Def)>::archetype_id(&reg);
        let metas: Vec<_> = <(Pos, Vel, Hp, Atk, Def)>::column_metas(&reg)
            .into_iter()
            .map(|cm| (cm.component_id, cm.meta))
            .collect();
        let mut arch = Archetype::new(aid, &metas, 4096);

        let e = Entity::new(0, 0);
        arch.add_entity_with_bundle(e, (Pos(1, 2), Vel(3, 4), Hp(100), Atk(10), Def(5)), &reg);
        assert_eq!(arch.len(), 1);
        assert_eq!(arch.get_comp::<Pos>(0, &reg), Some(&Pos(1, 2)));
        assert_eq!(arch.get_comp::<Hp>(0, &reg), Some(&Hp(100)));
        assert_eq!(arch.get_comp::<Def>(0, &reg), Some(&Def(5)));
    }

    #[test]
    fn archetype_builder_basic() {
        let reg = ComponentIdRegistry::new();
        let mut builder = ArchetypeBuilder::new();
        builder.add_component::<Pos>(&reg).add_component::<Vel>(&reg);

        let bundle_id = <(Pos, Vel)>::archetype_id(&reg);
        assert_eq!(builder.archetype_id(), bundle_id);

        let mut arch = builder.build(1024);
        let e = Entity::new(0, 0);
        let row = arch.push_entity(e);
        unsafe {
            arch.write_component(row, reg.id_for::<Pos>(), Pos(10, 20));
            arch.write_component(row, reg.id_for::<Vel>(), Vel(1, 2));
        }
        assert_eq!(arch.get_comp::<Pos>(0, &reg), Some(&Pos(10, 20)));
        assert_eq!(arch.get_comp::<Vel>(0, &reg), Some(&Vel(1, 2)));
    }

    #[test]
    fn archetype_builder_duplicate_ignored() {
        let reg = ComponentIdRegistry::new();
        let mut builder = ArchetypeBuilder::new();
        builder.add_component::<Pos>(&reg).add_component::<Pos>(&reg).add_component::<Vel>(&reg);
        assert_eq!(builder.metas().len(), 2);
    }

    #[test]
    fn archetype_builder_add_bundle() {
        let reg = ComponentIdRegistry::new();
        let mut builder = ArchetypeBuilder::new();
        builder.add_components::<(Pos, Vel)>(&reg);
        assert_eq!(builder.metas().len(), 2);

        let bundle_id = <(Pos, Vel)>::archetype_id(&reg);
        assert_eq!(builder.archetype_id(), bundle_id);
    }

    #[test]
    fn archetype_builder_add_bundle_with_overlap() {
        let reg = ComponentIdRegistry::new();
        let mut builder = ArchetypeBuilder::new();
        builder.add_component::<Pos>(&reg);
        builder.add_components::<(Pos, Vel, Hp)>(&reg);

        assert_eq!(builder.metas().len(), 3);

        let expected = <(Pos, Vel, Hp)>::archetype_id(&reg);
        assert_eq!(builder.archetype_id(), expected);
    }

    #[test]
    fn archetype_builder_add_multiple_bundles() {
        let reg = ComponentIdRegistry::new();
        let mut builder = ArchetypeBuilder::new();
        builder
            .add_components::<(Pos, Vel)>(&reg)
            .add_components::<(Hp, Atk)>(&reg);

        assert_eq!(builder.metas().len(), 4);

        let expected = <(Pos, Vel, Hp, Atk)>::archetype_id(&reg);
        assert_eq!(builder.archetype_id(), expected);
    }

    #[test]
    fn bitvec_beyond_256() {
        let mut id = ArchetypeId::empty();

        id.set(ComponentId(0));
        id.set(ComponentId(255));
        id.set(ComponentId(256));
        id.set(ComponentId(500));
        id.set(ComponentId(1000));

        assert!(id.has(ComponentId(0)));
        assert!(id.has(ComponentId(255)));
        assert!(id.has(ComponentId(256)));
        assert!(id.has(ComponentId(500)));
        assert!(id.has(ComponentId(1000)));
        assert!(!id.has(ComponentId(1)));
        assert!(!id.has(ComponentId(999)));
        assert!(!id.has(ComponentId(1001)));

        assert_eq!(id.count(), 5);

        let collected: Vec<_> = id.iter_set().collect();
        assert_eq!(collected, vec![
            ComponentId(0), ComponentId(255), ComponentId(256),
            ComponentId(500), ComponentId(1000),
        ]);
    }

    #[test]
    fn bitvec_contains_all_beyond_256() {
        let mut superset = ArchetypeId::empty();
        superset.set(ComponentId(0));
        superset.set(ComponentId(300));
        superset.set(ComponentId(500));

        let mut mask = ArchetypeId::empty();
        mask.set(ComponentId(0));
        mask.set(ComponentId(300));
        assert!(superset.contains_all(&mask));

        mask.set(ComponentId(999));
        assert!(!superset.contains_all(&mask));
    }

    #[test]
    fn bitvec_union_beyond_256() {
        let mut a = ArchetypeId::empty();
        a.set(ComponentId(0));
        a.set(ComponentId(300));

        let mut b = ArchetypeId::empty();
        b.set(ComponentId(100));
        b.set(ComponentId(500));

        let u = a.union(&b);
        assert!(u.has(ComponentId(0)));
        assert!(u.has(ComponentId(100)));
        assert!(u.has(ComponentId(300)));
        assert!(u.has(ComponentId(500)));
        assert_eq!(u.count(), 4);
    }

    #[test]
    fn bitvec_hash_consistency() {
        use std::collections::hash_map::DefaultHasher;

        let mut a = ArchetypeId::empty();
        a.set(ComponentId(300));

        let mut b = ArchetypeId::empty();
        b.set(ComponentId(300));

        let hash_a = {
            let mut h = DefaultHasher::new();
            a.hash(&mut h);
            h.finish()
        };
        let hash_b = {
            let mut h = DefaultHasher::new();
            b.hash(&mut h);
            h.finish()
        };
        assert_eq!(hash_a, hash_b);
        assert_eq!(a, b);
    }

    #[test]
    fn bitvec_empty_equality() {
        let a = ArchetypeId::empty();

        let mut b = ArchetypeId::empty();
        b.set(ComponentId(500));

        // b has a longer internal vec but unsetting the bit should yield an
        // equal id via the PartialEq impl (which compares SmallVec directly).
        // Since we don't have an unset op, just verify they're not equal.
        assert_ne!(a, b);
    }
}
