use alloc::vec;
use alloc::vec::Vec;
use core::num::NonZero;

#[derive(Copy, Clone, Debug)]
// Alignment repr necessary to allow LLVM to better output
#[repr(C, align(8))]
pub struct Entity {
    // Do not reorder the fields here. The ordering is explicitly used by repr(C)
    // to make this struct equivalent to a u64.
    #[cfg(target_endian = "little")]
    pub(crate) idx: u32,
    pub(crate) generation: NonZero<u32>,
    // Lol. Lmao. But here for Compiler Reasons™
    #[cfg(target_endian = "big")]
    pub(crate) idx: u32,
}

impl PartialEq for Entity {
    fn eq(&self, other: &Self) -> bool {
        self.to_bits().eq(&other.to_bits())
    }
}

impl Eq for Entity {}

impl core::hash::Hash for Entity {
    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
        self.to_bits().hash(state);
    }
}

impl core::fmt::Display for Entity {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        write!(f, "{}v{}", self.idx, self.generation.get())
    }
}

impl Entity {
    #[inline(always)]
    pub(crate) const fn to_bits(self) -> u64 {
        self.idx as u64 | (self.generation.get() as u64) << 32
    }
}

pub struct Census {
    storage: Vec<Option<NonZero<u32>>>,
    masks: Vec<u64>,
    lowest_open: usize,
}

impl Census {
    pub fn new() -> Self {
        Census {
            storage: vec![],
            masks: vec![],
            lowest_open: 0,
        }
    }

    #[inline(always)]
    unsafe fn inc(&mut self, idx: u32) -> NonZero<u32> {
        let generation = match unsafe { self.storage.get_unchecked(idx as usize) } {
            Some(generation) => generation.saturating_add(1),
            None => NonZero::<u32>::MIN,
        };
        self.storage[idx as usize] = Some(generation);
        generation
    }

    #[inline(always)]
    unsafe fn mark_empty(&mut self, idx: u32) {
        let mask_idx = (idx / u64::BITS) as usize;
        let mask_pos = (idx % u64::BITS) as u64;
        unsafe {
            *self.masks.get_unchecked_mut(mask_idx) |= 1u64 << mask_pos;
        }
        self.lowest_open = self.lowest_open.min(mask_idx);
    }

    #[inline(always)]
    pub fn reserve(&mut self) -> Entity {
        for (n, mask) in self
            .masks
            .iter()
            .enumerate()
            .skip(self.lowest_open)
            .filter(|(_, m)| **m != 0)
        {
            if let Some(mask_idx) = mask.lowest_one() {
                let idx = mask_idx + n as u32 * u64::BITS;
                self.lowest_open = self.lowest_open.max(n);
                // SAFTEY: idx/n come from iterator, must be smaller then storage
                unsafe {
                    *(self.masks.get_unchecked_mut(n)) -= 1u64 << mask_idx;
                    return Entity {
                        idx,
                        generation: self.inc(idx),
                    };
                }
            } else {
                unreachable!();
            }
        }

        let one = NonZero::<u32>::MIN;
        let idx = self.storage.len() as u32;
        self.storage.push(Some(one));
        self.storage.extend((0..63).map(|_| None));
        self.masks.push(u64::MAX - 1);
        self.lowest_open = self.masks.len() - 1;
        Entity {
            idx,
            generation: one,
        }
    }

    #[inline(always)]
    pub fn remove(&mut self, Entity { idx, .. }: Entity) {
        if idx as usize >= self.storage.len() {
            core::hint::cold_path();
            return;
        }

        // SAFETY: idx is smaller then size of storage
        unsafe {
            self.mark_empty(idx);
        }
    }

    #[inline(always)]
    pub fn validate(&self, Entity { idx, generation }: Entity) -> bool {
        if idx as usize >= self.storage.len() {
            core::hint::cold_path();
            return false;
        }
        unsafe {
            self.is_taken(idx)
                && (self
                    .storage
                    .get_unchecked(idx as usize)
                    .is_some_and(|g| g == generation))
        }
    }

    #[inline(always)]
    unsafe fn is_taken(&self, idx: u32) -> bool {
        let mask_idx = (idx / u64::BITS) as usize;
        let mask_pos = (idx % u64::BITS) as u64;
        unsafe { (*self.masks.get_unchecked(mask_idx) & 1u64 << mask_pos) == 0 }
    }
}

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

    #[test]
    fn reserve_validate() {
        let mut cen = Census::new();
        let e = cen.reserve();
        assert!(cen.validate(e));
        assert_eq!(e.idx, 0);
    }

    #[test]
    fn reserve_validate_large() {
        let mut cen = Census::new();
        let e1 = cen.reserve();
        for _ in 0..200 {
            cen.reserve();
        }
        let e2 = cen.reserve();

        assert!(cen.validate(e1));
        assert!(cen.validate(e2));
    }

    #[test]
    fn mask_backwards() {
        let mut cen = Census::new();
        let e1 = cen.reserve();
        for _ in 0..200 {
            cen.reserve();
        }
        cen.remove(e1);
        let e2 = cen.reserve();

        assert_eq!(e1.idx, e2.idx);
        assert_ne!(e1.generation, e2.generation);
    }

    #[test]
    fn reserve_despawn() {
        let mut cen = Census::new();
        let e = cen.reserve();
        assert!(cen.validate(e));
        cen.remove(e);
        assert!(!cen.validate(e));
    }

    #[test]
    fn reserve_two() {
        let mut cen = Census::new();
        let e1 = cen.reserve();
        let e2 = cen.reserve();

        assert_ne!(e1, e2);
        assert_ne!(e1.idx, e2.idx);
    }

    #[test]
    fn many() {
        let mut cen = Census::new();
        let first = cen.reserve();
        for _ in 0..128 {
            cen.reserve();
        }
        let last = cen.reserve();

        assert_ne!(last, first);
        assert_eq!(last.idx, 129);
    }

    #[test]
    fn recylce_simple() {
        let mut cen = Census::new();
        let e1 = cen.reserve();
        cen.remove(e1);
        let e2 = cen.reserve();
        assert_eq!(e1.idx, e2.idx);
        assert_ne!(e1.generation, e2.generation);
    }

    #[test]
    fn recycle_medium() {
        let mut cen = Census::new();
        let e1 = cen.reserve();
        let e2 = cen.reserve();
        let e3 = cen.reserve();

        cen.remove(e2);
        let e4 = cen.reserve();

        assert_eq!(e4.idx, e2.idx);
        assert_ne!(e4.generation, e2.generation);
        assert!(cen.validate(e1) && cen.validate(e3));
    }

    #[test]
    fn recylce_large() {
        let mut cen = Census::new();
        for _ in 0..200 {
            cen.reserve();
        }

        let e1 = cen.reserve();

        for _ in 0..200 {
            cen.reserve();
        }

        cen.remove(e1);
        let e2 = cen.reserve();

        assert_eq!(e1.idx, e2.idx);
        assert_ne!(e1.generation, e2.generation);
    }
}