use core::{num::NonZeroUsize, usize};

use allocator_api2::{
    alloc::{Allocator, Global},
    boxed::Box,
    vec::Vec,
};
use bytemuck::Contiguous;

const SPARSE_PAGESIZE: usize = (1 << 10) * 4;
type SparsePage<A> = Option<(Box<[Option<NonZeroUsize>; SPARSE_PAGESIZE], A>, usize)>;

pub struct SparseSet<T, PackedAlloc = Global, SparseAlloc = Global>
where
    PackedAlloc: Allocator,
    SparseAlloc: Allocator,
{
    sparse: Vec<SparsePage<SparseAlloc>, SparseAlloc>,
    dense: Vec<T, PackedAlloc>,
    dense_to_id: Vec<usize, SparseAlloc>,
}

impl<T> SparseSet<T> {
    pub const fn new() -> Self {
        Self {
            sparse: Vec::new(),
            dense: Vec::new(),
            dense_to_id: Vec::new(),
        }
    }
}

impl<T, PackedAlloc, SparseAlloc> SparseSet<T, PackedAlloc, SparseAlloc>
where
    PackedAlloc: Allocator,
    SparseAlloc: Allocator + Clone,
{
    pub fn insert(&mut self, id: usize, value: T) -> Option<T> {
        match self.get_dense_idx(id) {
            Some(idx) => {
                let previous = core::mem::replace(&mut self.dense[idx], value);
                self.dense_to_id[idx] = id;
                Some(previous)
            }
            None => {
                self.increase_page_usage_count(id);
                self.set_dense_idx(id, Some(self.dense.len()));
                self.dense.push(value);
                self.dense_to_id.push(id);
                None
            }
        }
    }

    pub fn get(&self, id: usize) -> Option<&T> {
        self.dense.get(self.get_dense_idx(id)?)
    }

    pub fn get_mut(&mut self, id: usize) -> Option<&mut T> {
        let idx = self.get_dense_idx(id)?;
        self.dense.get_mut(idx)
    }

    fn set_dense_idx(&mut self, id: usize, idx: Option<usize>) {
        let page = id / SPARSE_PAGESIZE;
        let sparse_index = id % SPARSE_PAGESIZE;

        if page >= self.sparse.len() {
            self.sparse.resize(page + 1, None);
        }

        if self.sparse[page].is_none() {
            self.sparse[page] = Some((
                Box::new_in([None; 4096], self.sparse.allocator().clone()),
                1,
            ))
        }

        match &mut self.sparse[page] {
            Some(page) => {
                page.0[sparse_index] = idx.map(|i| NonZeroUsize::new(i + 1).unwrap());
            }
            None => unreachable!("wtf, failed to init sparse page 5 lines above??"),
        }
    }

    pub fn get_dense_idx(&self, id: usize) -> Option<usize> {
        let page = id / SPARSE_PAGESIZE;
        let sparse_index = id % SPARSE_PAGESIZE;
        let page = self.sparse.get(page)?.as_ref()?;
        page.0[sparse_index].map(|idx| idx.into_integer() - 1)
    }

    fn reduce_page_usage_count(&mut self, id: usize) {
        let page = id / SPARSE_PAGESIZE;
        let Some(usage) = &mut self.sparse[page] else {
            return;
        };
        usage.1 -= 1;
        let usage = usage.1;
        if usage == 0 {
            self.sparse[page] = None;
        }
    }

    fn increase_page_usage_count(&mut self, id: usize) {
        let page = id / SPARSE_PAGESIZE;
        if page >= self.sparse.len() {
            return;
        }
        let Some(usage) = &mut self.sparse[page] else {
            return;
        };
        usage.1 += 1;
    }

    pub fn remove(&mut self, id: usize) -> Option<T> {
        let index = self.get_dense_idx(id)?;
        if self.dense.is_empty() {
            return None;
        }

        self.set_dense_idx(*self.dense_to_id.last().unwrap(), Some(index));
        self.set_dense_idx(id, None);
        self.reduce_page_usage_count(id);

        let previous = self.dense.swap_remove(index);
        self.dense_to_id.swap_remove(index);
        Some(previous)
    }

    pub fn is_emtpy(&self) -> bool {
        self.len() == 0
    }

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

    pub fn contains(&self, id: usize) -> bool {
        self.get_dense_idx(id).is_some()
    }

    pub fn keys(&self) -> &[usize] {
        &self.dense_to_id
    }

    pub fn values(&self) -> &[T] {
        &self.dense
    }

    pub fn new_in(packed_alloc: PackedAlloc, sparse_alloc: SparseAlloc) -> Self {
        Self {
            dense: Vec::new_in(packed_alloc),
            sparse: Vec::new_in(sparse_alloc.clone()),
            dense_to_id: Vec::new_in(sparse_alloc),
        }
    }
}

impl<T, PackedAlloc> SparseSet<T, PackedAlloc>
where
    PackedAlloc: Allocator,
{
    pub const fn new_in_packed(packed_alloc: PackedAlloc) -> Self {
        Self {
            sparse: Vec::new(),
            dense: Vec::new_in(packed_alloc),
            dense_to_id: Vec::new(),
        }
    }
}

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

    #[test]
    fn insert() {
        let mut sparse_set = SparseSet::<u32>::new();
        sparse_set.insert(10, 1);
        assert_eq!(sparse_set.keys(), &[10]);
        assert_eq!(sparse_set.values(), &[1]);

        assert_eq!(
            sparse_set.sparse[0].as_ref().unwrap().0[10].unwrap(),
            NonZeroUsize::new(1).unwrap()
        );
        assert_eq!(sparse_set.sparse[0].as_ref().unwrap().1, 1);
        assert_eq!(sparse_set.insert(10, 2).unwrap(), 1);
        assert_eq!(sparse_set.values(), &[2]);
        assert_eq!(sparse_set.sparse[0].as_ref().unwrap().1, 1);

        sparse_set.insert(11, 4);
        assert_eq!(sparse_set.keys(), &[10, 11]);
        assert_eq!(sparse_set.values(), &[2, 4]);
        assert_eq!(
            sparse_set.sparse[0].as_ref().unwrap().0[11].unwrap(),
            NonZeroUsize::new(2).unwrap()
        );
        assert_eq!(sparse_set.sparse[0].as_ref().unwrap().1, 2);

        sparse_set.insert(5000, 3);
        assert_eq!(sparse_set.keys(), &[10, 11, 5000]);
        assert_eq!(sparse_set.values(), &[2, 4, 3]);
        assert_eq!(
            sparse_set.sparse[5000 / SPARSE_PAGESIZE]
                .as_ref()
                .unwrap()
                .0[5000 % SPARSE_PAGESIZE]
                .unwrap(),
            NonZeroUsize::new(3).unwrap()
        );
        assert_eq!(
            sparse_set.sparse[5000 / SPARSE_PAGESIZE]
                .as_ref()
                .unwrap()
                .1,
            1
        );

        assert_eq!(*sparse_set.get(10).unwrap(), 2);
        assert_eq!(*sparse_set.get(11).unwrap(), 4);
        assert_eq!(*sparse_set.get(5000).unwrap(), 3);
    }

    #[test]
    fn remove() {
        let mut sparse_set = SparseSet::<u32>::new();
        sparse_set.insert(10, 1);
        sparse_set.insert(11, 2);
        sparse_set.insert(12, 2);
        sparse_set.insert(SPARSE_PAGESIZE, 1);
        sparse_set.insert(SPARSE_PAGESIZE + 1, 2);
        sparse_set.insert(SPARSE_PAGESIZE + 2, 3);

        assert_eq!(sparse_set.remove(SPARSE_PAGESIZE + 2).unwrap(), 3);
        assert_eq!(sparse_set.sparse[1].as_ref().unwrap().1, 2);
        assert_eq!(sparse_set.keys(), [10, 11, 12, SPARSE_PAGESIZE, SPARSE_PAGESIZE + 1]);
        assert_eq!(sparse_set.values(), [1, 2, 2, 1, 2]);

        assert_eq!(sparse_set.remove(SPARSE_PAGESIZE + 1).unwrap(), 2);
        assert_eq!(sparse_set.sparse[1].as_ref().unwrap().1, 1);
        assert_eq!(sparse_set.keys(), [10, 11, 12, SPARSE_PAGESIZE]);
        assert_eq!(sparse_set.values(), [1, 2, 2, 1]);

        assert_eq!(sparse_set.remove(SPARSE_PAGESIZE).unwrap(), 1);
        assert!(sparse_set.sparse[1].is_none());
        assert_eq!(sparse_set.keys(), [10, 11, 12]);
        assert_eq!(sparse_set.values(), [1, 2, 2]);


        sparse_set.insert(SPARSE_PAGESIZE, 1);
        sparse_set.insert(SPARSE_PAGESIZE + 1, 2);
        sparse_set.insert(SPARSE_PAGESIZE + 2, 3);

        assert_eq!(sparse_set.remove(10).unwrap(), 1);
        assert_eq!(sparse_set.sparse[0].as_ref().unwrap().1, 2);
        // swap-remove 
        assert_eq!(sparse_set.keys(), [SPARSE_PAGESIZE + 2, 11, 12, SPARSE_PAGESIZE, SPARSE_PAGESIZE + 1]);
        assert_eq!(sparse_set.values(), [3, 2, 2, 1, 2]);

        assert_eq!(sparse_set.remove(11).unwrap(), 2);
        assert_eq!(sparse_set.sparse[0].as_ref().unwrap().1, 1);
        assert_eq!(sparse_set.keys(), [SPARSE_PAGESIZE + 2, SPARSE_PAGESIZE + 1, 12, SPARSE_PAGESIZE]);
        assert_eq!(sparse_set.values(), [3, 2, 2, 1]);

        assert_eq!(sparse_set.remove(12).unwrap(), 2);
        assert!(sparse_set.sparse[0].is_none());
        assert_eq!(sparse_set.keys(), [SPARSE_PAGESIZE + 2, SPARSE_PAGESIZE + 1, SPARSE_PAGESIZE]);
        assert_eq!(sparse_set.values(), [3, 2, 1]);
    }
}