armour 0.30.27

use std::marker::PhantomData;

use crate::{GetType, InlineArray, Typ};
use bytemuck::{Zeroable, cast_ref};
use derive_more::{AsMut, AsRef, Deref, DerefMut};

/// indexed collection of items for ids with inner u64
#[derive(Debug, Clone, AsRef, Deref, DerefMut, AsMut)]
#[repr(transparent)]
pub struct IndexedCollection<T> {
    #[deref]
    #[deref_mut]
    data: InlineArray,
    _phantom: std::marker::PhantomData<T>,
}

impl<T> From<InlineArray> for IndexedCollection<T> {
    fn from(inline: InlineArray) -> Self {
        Self {
            data: inline,
            _phantom: PhantomData,
        }
    }
}

impl<T> From<IndexedCollection<T>> for InlineArray {
    fn from(collection: IndexedCollection<T>) -> Self {
        collection.data
    }
}

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

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

    pub fn count(&self) -> u64 {
        let header_bytes = self
            .data
            .as_ref()
            .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.as_ref().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.make_mut());
        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 = InlineArray::with_len(bytes_len * 2);
            let new_slice = new.make_mut();
            new_slice[..bytes_len].copy_from_slice(data.as_ref());
            let mut view = HeaderPrefixedArray::from_mut_slice(new_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.make_mut());
        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.make_mut().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.as_ref().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 data = self.data.make_mut();
        let (header_bytes, items_bytes) = 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.make_mut())
    }

    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 = InlineArray::with_len(old_len * mul);
        new_data.make_mut()[..old_len].copy_from_slice(self.data.as_ref());
        self.data = new_data;
    }
}

impl<T> GetType for IndexedCollection<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
        }
    }
}