use std::alloc::{alloc, handle_alloc_error, realloc, Layout};
use std::cmp::max;
use std::marker::PhantomData;
use std::mem::{forget, needs_drop};
use std::ptr::NonNull;

/// A list of `ItemT`. This data structure stores a list for every field of `ItemT`,
/// reducing memory usage if `ItemT` contains padding and improves memory cache usage if
/// only certain fields are needed when iterating.
///
/// Inspired by Zig's `MultiArrayList`.
///
/// Note: All of the lists are stored in the same allocation.
///
/// For example, if you have three of the following struct:
/// ```
/// struct Person
/// {
///     first_name: String,
///     age: u8,
/// }
/// ```
///
/// It would be stored like this in memory:
/// ```text
/// first_name, first_name, first_name,
/// age, age, age,
/// ```
#[derive(Debug)]
pub struct MultiVec<ItemT>
where
    ItemT: Item,
{
    _pd: PhantomData<ItemT>,
    ptr: NonNull<u8>,
    field_arr_byte_offsets: Vec<usize>,
    length: usize,
    capacity: usize,
    layout: Option<Layout>,
}

impl<ItemT> MultiVec<ItemT>
where
    ItemT: Item,
{
    // The following is borrow from std's RawVec implementation:
    // Skip to:
    // - 8 if the element size is 1, because any heap allocators is likely to round up a
    //   request of less than 8 bytes to at least 8 bytes.
    // - 4 if elements are moderate-sized (<= 1 KiB).
    // - 1 otherwise, to avoid wasting too much space for very short Vecs.
    const MIN_NON_ZERO_CAP: usize = if size_of::<ItemT>() == 1 {
        8
    } else if size_of::<ItemT>() <= 1024 {
        4
    } else {
        1
    };

    /// Returns a new `MultiVec`. This function does not allocate any memory.
    pub const fn new() -> Self
    {
        Self {
            _pd: PhantomData,
            ptr: NonNull::dangling(),
            field_arr_byte_offsets: Vec::new(),
            length: 0,
            capacity: 0,
            layout: None,
        }
    }

    pub fn with_capacity(capacity: usize) -> Self
    {
        let mut this = Self {
            _pd: PhantomData,
            ptr: NonNull::dangling(),
            field_arr_byte_offsets: Vec::new(),
            length: 0,
            capacity: 0,
            layout: None,
        };

        this.do_first_alloc(capacity);

        this
    }

    /// Pushes a item to the `MultiVec`.
    ///
    /// ## Note on performance
    /// Pushing can be pretty slow. Since all of the field lists are stored in the same
    /// allocation, when pushing and the `MultiVec` needs to grow, all lists except the
    /// first has to be moved to new locations for them to not overlap.
    pub fn push(&mut self, item: ItemT)
    {
        if self.capacity != 0 {
            if self.capacity == self.length {
                self.grow_amortized(1);
            }

            self.write_item(self.length - 1, item);

            self.length += 1;

            return;
        }

        self.do_first_alloc(1);

        self.write_item(0, item);

        self.length = 1;
    }

    /// Returns a field of the item with the given index.
    pub fn get<FieldSel>(
        &self,
        index: usize,
    ) -> Option<&<FieldSel as ItemFieldSelection<ItemT>>::Field>
    where
        FieldSel: ItemFieldSelection<ItemT>,
    {
        if index >= self.length {
            return None;
        }

        let field_metadata = FieldSel::metadata();

        let field_arr_byte_offset = self.field_arr_byte_offsets[FieldSel::INDEX];

        let field_arr_ptr = unsafe { self.ptr.byte_add(field_arr_byte_offset) };

        let field_ptr = unsafe { field_arr_ptr.add(field_metadata.size * index) };

        Some(unsafe { field_ptr.cast().as_ref() })
    }

    /// Returns the number of items stored in this `MultiVec`.
    pub fn len(&self) -> usize
    {
        self.length
    }

    /// Returns how many items this `MultiVec` has capacity for.
    pub fn capacity(&self) -> usize
    {
        self.capacity
    }

    /// Returns whether this `MultiVec` is empty.
    pub fn is_empty(&self) -> bool
    {
        self.length == 0
    }

    fn grow_amortized(&mut self, additional: usize)
    {
        let required_cap = self.capacity.checked_add(additional).unwrap();

        // This guarantees exponential growth. The doubling cannot overflow
        // because `cap <= isize::MAX` and the type of `cap` is `usize`.
        let new_capacity = max(self.capacity * 2, required_cap);
        let new_capacity = max(Self::MIN_NON_ZERO_CAP, new_capacity);

        let layout = &self.layout.unwrap();

        let (new_layout, new_field_arr_byte_offsets) = Self::create_layout(new_capacity);

        let Some(new_ptr) = NonNull::new(unsafe {
            realloc(self.ptr.as_ptr(), *layout, new_layout.size())
        }) else {
            handle_alloc_error(new_layout);
        };

        for (field_index, field_metadata) in
            ItemT::iter_field_metadata().enumerate().rev()
        {
            let old_field_arr_byte_offset = self.field_arr_byte_offsets[field_index];
            let new_field_arr_byte_offset = new_field_arr_byte_offsets[field_index];

            let old_field_arr_ptr =
                unsafe { new_ptr.byte_add(old_field_arr_byte_offset) };

            let new_field_arr_ptr =
                unsafe { new_ptr.byte_add(new_field_arr_byte_offset) };

            unsafe {
                std::ptr::copy(
                    old_field_arr_ptr.as_ptr(),
                    new_field_arr_ptr.as_ptr(),
                    field_metadata.size * self.capacity,
                );
            }
        }

        self.ptr = new_ptr;
        self.layout = Some(new_layout);
        self.capacity = new_capacity;
        self.field_arr_byte_offsets = new_field_arr_byte_offsets;
    }

    fn do_first_alloc(&mut self, capacity: usize)
    {
        let (layout, field_arr_byte_offsets) = Self::create_layout(capacity);

        let Some(ptr) = NonNull::new(unsafe { alloc(layout) }) else {
            handle_alloc_error(layout);
        };

        self.ptr = ptr;
        self.capacity = capacity;
        self.field_arr_byte_offsets = field_arr_byte_offsets;
        self.layout = Some(layout);
    }

    fn create_layout(length: usize) -> (Layout, Vec<usize>)
    {
        let mut field_metadata_iter = ItemT::iter_field_metadata();

        let first_field_metadata = field_metadata_iter.next().unwrap();

        let mut layout = array_layout(
            first_field_metadata.size,
            first_field_metadata.alignment,
            length,
        )
        .unwrap();

        let mut field_arr_byte_offsets = Vec::with_capacity(ItemT::FIELD_CNT);

        field_arr_byte_offsets.push(0);

        for field_metadata in field_metadata_iter {
            let (new_layout, array_byte_offset) = layout
                .extend(
                    array_layout(field_metadata.size, field_metadata.alignment, length)
                        .unwrap(),
                )
                .unwrap();

            layout = new_layout;

            field_arr_byte_offsets.push(array_byte_offset);
        }

        (layout, field_arr_byte_offsets)
    }

    fn write_item(&mut self, index: usize, item: ItemT)
    {
        for (field_index, field_metadata) in ItemT::iter_field_metadata().enumerate() {
            let field_arr_byte_offset = self.field_arr_byte_offsets[field_index];

            let field_arr_ptr = unsafe { self.ptr.byte_add(field_arr_byte_offset) };

            let field_ptr = unsafe { field_arr_ptr.add(field_metadata.size * index) };

            unsafe {
                std::ptr::copy_nonoverlapping(
                    (&item as *const ItemT).byte_add(field_metadata.offset) as *const u8,
                    field_ptr.as_ptr(),
                    field_metadata.size,
                );
            }
        }

        forget(item);
    }
}

impl<ItemT> Drop for MultiVec<ItemT>
where
    ItemT: Item,
{
    fn drop(&mut self)
    {
        if needs_drop::<ItemT>() {
            for index in 0..self.length {
                for (field_index, field_metadata) in
                    ItemT::iter_field_metadata().enumerate()
                {
                    let field_arr_byte_offset = self.field_arr_byte_offsets[field_index];

                    let field_arr_ptr =
                        unsafe { self.ptr.byte_add(field_arr_byte_offset) };

                    let field_ptr =
                        unsafe { field_arr_ptr.add(field_metadata.size * index) };

                    unsafe {
                        ItemT::drop_field_inplace(field_index, field_ptr.as_ptr());
                    }
                }
            }
        }

        if let Some(layout) = self.layout {
            unsafe {
                std::alloc::dealloc(self.ptr.as_ptr(), layout);
            }
        }
    }
}

/// Usable as a item of a [`MultiVec`].
///
/// # Safety
/// The iterator returned by `iter_field_metadata` must yield [`ItemFieldMetadata`] that
/// correctly represents fields of the implementor type.
pub unsafe trait Item
{
    type FieldMetadataIter<'a>: Iterator<Item = &'a ItemFieldMetadata>
        + DoubleEndedIterator
        + ExactSizeIterator;

    /// The number of fields this item has.
    const FIELD_CNT: usize;

    /// Returns a iterator of metadata of the fields of this item.
    fn iter_field_metadata() -> Self::FieldMetadataIter<'static>;

    /// Drops a field of this item inplace.
    ///
    /// # Safety
    /// Behavior is undefined if any of the following conditions are violated:
    ///
    /// * `field_index` must be a valid field index for this item.
    /// * The value `field_ptr` points to must be valid for the type of the field with
    ///   index `field_index`.
    /// * `field_ptr` must be [valid] for both reads and writes.
    /// * `field_ptr` must be properly aligned, even if the field type has size 0.
    /// * `field_ptr` must be nonnull, even if the field type has size 0.
    /// * The value `field_ptr` points to must be valid for dropping, which may mean it
    ///   must uphold additional invariants. These invariants depend on the type of the
    ///   value being dropped. For instance, when dropping a Box, the box's pointer to the
    ///   heap must be valid.
    /// * While `drop_field_inplace` is executing, the only way to access parts of
    ///   `field_ptr` is through the `&mut self` references supplied to the [`Drop::drop`]
    ///   methods that `drop_field_inplace` invokes.
    ///
    /// Additionally, if the field type is not [`Copy`], using the pointed-to value after
    /// calling `drop_field_inplace` can cause undefined behavior. Note that `*field_ptr =
    /// foo` counts as a use because it will cause the value to be dropped again.
    ///
    /// [valid]: https://doc.rust-lang.org/std/ptr/index.html#safety
    unsafe fn drop_field_inplace(field_index: usize, field_ptr: *mut u8);
}

/// Contains metadata of a field of a [`Item`].
///
/// # Examples
/// ```
/// # use multi_vec::ItemFieldMetadata;
/// #
/// struct CatFood
/// {
///     label: String,
///     age_days: u16,
///     quality: u8,
/// }
///
/// let cat_food_field_metadata = ItemFieldMetadata {
///     offset: std::mem::offset_of!(CatFood, age_days),
///     size: std::mem::size_of::<u16>(),
///     alignment: std::mem::align_of::<u16>(),
/// };
/// ```
pub struct ItemFieldMetadata
{
    pub offset: usize,
    pub size: usize,
    pub alignment: usize,
}

/// A field selection for `ItemT`.
///
/// # Safety
/// The constant `INDEX`, the type `Field` and the `ItemFieldMetadata` returned by the
/// `metadata` function must correctly represent a field of `ItemT`;
pub unsafe trait ItemFieldSelection<ItemT>
where
    ItemT: Item,
{
    const INDEX: usize;

    type Field;

    /// Returns metadata of this item field.
    fn metadata() -> &'static ItemFieldMetadata;
}

#[inline]
const fn array_layout(
    element_size: usize,
    align: usize,
    n: usize,
) -> Result<Layout, CoolLayoutError>
{
    // We need to check two things about the size:
    //  - That the total size won't overflow a `usize`, and
    //  - That the total size still fits in an `isize`.
    // By using division we can check them both with a single threshold.
    // That'd usually be a bad idea, but thankfully here the element size
    // and alignment are constants, so the compiler will fold all of it.
    if element_size != 0 && n > max_size_for_align(align) / element_size {
        return Err(CoolLayoutError);
    }

    // SAFETY: We just checked that we won't overflow `usize` when we multiply.
    // This is a useless hint inside this function, but after inlining this helps
    // deduplicate checks for whether the overall capacity is zero (e.g., in RawVec's
    // allocation path) before/after this multiplication.
    let array_size = unsafe { element_size.unchecked_mul(n) };

    // SAFETY: We just checked above that the `array_size` will not
    // exceed `isize::MAX` even when rounded up to the alignment.
    // And `Alignment` guarantees it's a power of two.
    unsafe { Ok(Layout::from_size_align_unchecked(array_size, align)) }
}

#[inline(always)]
const fn max_size_for_align(align: usize) -> usize
{
    // (power-of-two implies align != 0.)

    // Rounded up size is:
    //   size_rounded_up = (size + align - 1) & !(align - 1);
    //
    // We know from above that align != 0. If adding (align - 1)
    // does not overflow, then rounding up will be fine.
    //
    // Conversely, &-masking with !(align - 1) will subtract off
    // only low-order-bits. Thus if overflow occurs with the sum,
    // the &-mask cannot subtract enough to undo that overflow.
    //
    // Above implies that checking for summation overflow is both
    // necessary and sufficient.
    isize::MAX as usize - (align - 1)
}

#[derive(Debug)]
struct CoolLayoutError;

#[cfg(test)]
mod tests
{
    use std::mem::offset_of;
    use std::ptr::drop_in_place;

    use crate::{Item, ItemFieldMetadata, ItemFieldSelection, MultiVec};

    struct Foo
    {
        num_a: u32,
        num_b: u16,
    }

    struct FooFieldNumA;
    struct FooFieldNumB;

    unsafe impl ItemFieldSelection<Foo> for FooFieldNumA
    {
        type Field = u32;

        const INDEX: usize = 0;

        fn metadata() -> &'static ItemFieldMetadata
        {
            &FOO_FIELD_METADATA[0]
        }
    }

    unsafe impl ItemFieldSelection<Foo> for FooFieldNumB
    {
        type Field = u16;

        const INDEX: usize = 1;

        fn metadata() -> &'static ItemFieldMetadata
        {
            &FOO_FIELD_METADATA[1]
        }
    }

    static FOO_FIELD_METADATA: [ItemFieldMetadata; 2] = [
        ItemFieldMetadata {
            offset: offset_of!(Foo, num_a),
            size: size_of::<u32>(),
            alignment: align_of::<u32>(),
        },
        ItemFieldMetadata {
            offset: offset_of!(Foo, num_b),
            size: size_of::<u16>(),
            alignment: align_of::<u16>(),
        },
    ];

    unsafe impl Item for Foo
    {
        type FieldMetadataIter<'a> = std::slice::Iter<'a, ItemFieldMetadata>;

        const FIELD_CNT: usize = 2;

        fn iter_field_metadata() -> Self::FieldMetadataIter<'static>
        {
            FOO_FIELD_METADATA.iter()
        }

        unsafe fn drop_field_inplace(field_index: usize, field_ptr: *mut u8)
        {
            if field_index == 0 {
                unsafe { drop_in_place::<u32>(field_ptr.cast()) }
            } else if field_index == 1 {
                unsafe { drop_in_place::<u16>(field_ptr.cast()) }
            }
        }
    }

    #[test]
    fn single_push_works()
    {
        let mut multi_vec = MultiVec::<Foo>::new();

        multi_vec.push(Foo { num_a: 123, num_b: 654 });

        assert_eq!(multi_vec.capacity, 1);
        assert_eq!(multi_vec.length, 1);

        assert_eq!(multi_vec.field_arr_byte_offsets, [0, size_of::<u32>()]);

        assert_eq!(
            unsafe {
                std::slice::from_raw_parts::<u32>(multi_vec.ptr.as_ptr().cast(), 1)
            },
            [123]
        );

        assert_eq!(
            unsafe {
                std::slice::from_raw_parts::<u16>(
                    multi_vec.ptr.as_ptr().byte_add(size_of::<u32>()).cast(),
                    1,
                )
            },
            [654]
        );
    }

    #[test]
    fn multiple_pushes_works()
    {
        let mut multi_vec = MultiVec::<Foo>::new();

        multi_vec.push(Foo { num_a: u32::MAX / 2, num_b: 654 });
        multi_vec.push(Foo { num_a: 765, num_b: u16::MAX / 3 });
        multi_vec.push(Foo { num_a: u32::MAX / 5, num_b: 337 });

        assert_eq!(multi_vec.capacity, 4);
        assert_eq!(multi_vec.length, 3);

        assert_eq!(multi_vec.field_arr_byte_offsets, [0, size_of::<u32>() * 4]);

        assert_eq!(
            unsafe {
                std::slice::from_raw_parts::<u32>(multi_vec.ptr.as_ptr().cast(), 3)
            },
            [u32::MAX / 2, 765, u32::MAX / 5]
        );

        assert_eq!(
            unsafe {
                std::slice::from_raw_parts::<u16>(
                    multi_vec.ptr.as_ptr().byte_add(size_of::<u32>() * 4).cast(),
                    3,
                )
            },
            [654, u16::MAX / 3, 337]
        );
    }

    #[test]
    fn multiple_pushes_in_preallocated_works()
    {
        let mut multi_vec = MultiVec::<Foo>::with_capacity(2);

        multi_vec.push(Foo { num_a: 83710000, num_b: 654 });
        multi_vec.push(Foo { num_a: 765, num_b: u16::MAX / 7 });

        assert_eq!(multi_vec.capacity, 2);
        assert_eq!(multi_vec.length, 2);

        assert_eq!(multi_vec.field_arr_byte_offsets, [0, size_of::<u32>() * 2]);

        assert_eq!(
            unsafe {
                std::slice::from_raw_parts::<u32>(multi_vec.ptr.as_ptr().cast(), 2)
            },
            [83710000, 765]
        );

        assert_eq!(
            unsafe {
                std::slice::from_raw_parts::<u16>(
                    multi_vec.ptr.as_ptr().byte_add(size_of::<u32>() * 2).cast(),
                    2,
                )
            },
            [654, u16::MAX / 7]
        );
    }
}