armour 0.30.27

use std::marker::PhantomData;

use crate::{GetType, Typ};
use bytemuck::{Zeroable, cast_ref};

/// Compact storage for a collection of `bytemuck::Pod` identifiers backed by `Vec<u8>`.
///
/// Layout: `[count: u64 LE][item_0: T][item_1: T]…[padding…]`
///
/// Uses `Vec<u8>` as the backing buffer, which guarantees heap-alignment
/// sufficient for `bytemuck::cast_slice` on common allocators.
/// This makes `IdVec` compatible with storage backends that do not provide
/// alignment guarantees (e.g. `fjall::Slice`): simply copy into a `Vec<u8>`
/// and wrap with `IdVec::from(vec)`.
///
/// The buffer grows by doubling when capacity is exhausted (`push`/`grow`).
/// Removal is O(1) via swap-remove with the last element.
#[derive(Debug, Clone)]
pub struct IdVec<T> {
    data: Vec<u8>,
    _phantom: std::marker::PhantomData<T>,
}

impl<T> From<Vec<u8>> for IdVec<T> {
    fn from(data: Vec<u8>) -> Self {
        Self {
            data,
            _phantom: PhantomData,
        }
    }
}

impl<T> From<IdVec<T>> for Vec<u8> {
    fn from(collection: IdVec<T>) -> Self {
        collection.data
    }
}

const SIZE: usize = std::mem::size_of::<u64>();

impl<T: bytemuck::Pod + Zeroable> IdVec<T> {
    pub fn new(item: &T) -> Self {
        let mut data = vec![0u8; SIZE * 4];
        data[..SIZE].copy_from_slice(&1u64.to_le_bytes());
        let item: &[u8; SIZE] = cast_ref(item);
        data[SIZE..(SIZE * 2)].copy_from_slice(item);
        Self {
            data,
            _phantom: PhantomData,
        }
    }

    pub fn count(&self) -> u64 {
        let header_bytes = self.data.get(..SIZE).expect("Header bytes not found");
        *bytemuck::from_bytes::<u64>(header_bytes)
    }

    #[doc(alias = "items")]
    pub fn contains(&self, item: Option<&T>) -> (u64, bool)
    where
        T: Eq,
    {
        let (header_bytes, items_bytes) = self.data.split_at(SIZE);
        let len = *bytemuck::from_bytes::<u64>(header_bytes);
        let slice: &[T] = bytemuck::cast_slice(items_bytes);
        let slice = &slice[..len as usize];
        let contains = item.is_some_and(|id| slice.contains(id));
        (len, contains)
    }

    /// return self, len, success
    pub fn push(self, item: &T, check: bool) -> (Self, u64, bool)
    where
        T: Eq,
    {
        let mut data = self.data;
        let mut view = HeaderPrefixedArray::from_mut_slice(data.as_mut_slice());
        if check && view.as_mut_slice().contains(item) {
            let len = *view.header;
            return (
                Self {
                    data,
                    _phantom: PhantomData,
                },
                len,
                false,
            );
        }
        if let Some(len) = view.push(*item) {
            (
                Self {
                    data,
                    _phantom: PhantomData,
                },
                len,
                true,
            )
        } else {
            let bytes_len = data.len();
            let mut new_data = vec![0u8; bytes_len * 2];
            new_data[..bytes_len].copy_from_slice(data.as_slice());
            let mut view = HeaderPrefixedArray::from_mut_slice(new_data.as_mut_slice());
            let len = view.push(*item).expect("Failed to push item");
            (
                Self {
                    data,
                    _phantom: PhantomData,
                },
                len,
                true,
            )
        }
    }

    /// return self, len, success
    pub fn remove(self, item: &T) -> (Self, u64, bool)
    where
        T: Eq,
    {
        let mut data = self.data;
        let mut view = HeaderPrefixedArray::from_mut_slice(data.as_mut_slice());
        let len = *view.header;
        let slice = view.as_mut_slice();
        let index = slice.iter().rposition(|x| x == item);
        match index {
            Some(index) => {
                if len == 1 {
                    data.fill(0);
                    return (
                        Self {
                            data,
                            _phantom: PhantomData,
                        },
                        0,
                        true,
                    );
                }
                let last_index = len - 1;
                slice.swap(index, last_index as usize);
                slice[last_index as usize] = T::zeroed();
                *view.header = last_index;
                (
                    Self {
                        data,
                        _phantom: PhantomData,
                    },
                    last_index,
                    true,
                )
            }
            None => (
                Self {
                    data,
                    _phantom: PhantomData,
                },
                len,
                false,
            ),
        }
    }

    /// returns a reference to the header and items (may be empty/uninitialized)
    pub fn split_data(&self) -> (u64, &[T]) {
        let (header_bytes, items_bytes) = self.data.split_at(SIZE);
        let count = *bytemuck::from_bytes::<u64>(header_bytes);
        let items = bytemuck::cast_slice(items_bytes);
        (count, items)
    }

    /// returns a mutable reference to the header and items (may be empty/uninitialized)
    pub fn split_data_mut(&mut self) -> (&mut u64, &mut [T]) {
        let (header_bytes, items_bytes) = self.data.split_at_mut(SIZE);
        let count = bytemuck::from_bytes_mut::<u64>(header_bytes);
        let items = bytemuck::cast_slice_mut(items_bytes);
        (count, items)
    }

    pub fn as_array(&mut self) -> HeaderPrefixedArray<'_, T> {
        HeaderPrefixedArray::from_mut_slice(self.data.as_mut_slice())
    }

    pub fn capacity(&self) -> usize {
        (self.data.len() - SIZE) / SIZE
    }

    pub fn grow(&mut self, mul: usize) {
        let old_len = self.data.len();
        let mut new_data = vec![0u8; old_len * mul];
        new_data[..old_len].copy_from_slice(self.data.as_slice());
        self.data = new_data;
    }
}

impl<T> GetType for IdVec<T> {
    const TYPE: Typ = Typ::Custom("indexed", &[Typ::U64, Typ::Vec(&Typ::Id64)]);
}

pub struct HeaderPrefixedArray<'a, T> {
    header: &'a mut u64,
    items: &'a mut [T],
}

impl<'a, T: bytemuck::Pod + Zeroable> HeaderPrefixedArray<'a, T> {
    pub fn from_mut_slice(slice: &'a mut [u8]) -> Self {
        let (header_bytes, items_bytes) = slice.split_at_mut(SIZE);
        let header = bytemuck::from_bytes_mut::<u64>(header_bytes);

        let items: &mut [T] = bytemuck::cast_slice_mut(items_bytes);

        Self { header, items }
    }

    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        let len = *self.header as usize;
        &mut self.items[..len]
    }

    pub fn push(&mut self, item: T) -> Option<u64> {
        let len = *self.header as usize;
        if len < self.items.len() {
            self.items[len] = item;
            let new_len = (len + 1) as u64;
            *self.header = new_len;
            Some(new_len)
        } else {
            None
        }
    }
}