use core::marker::PhantomData;

/// Trait that says that T can be safely transmuted into F
pub unsafe trait TransmutableFrom<T: Sized>: Sized {
    fn transmute_from(val: T) -> Self;

    fn transmute_slice_from(val: &[T]) -> &[Self] {
        let len = val.len();
        let ptr = val.as_ptr() as *const Self;

        unsafe { core::slice::from_raw_parts(ptr, len) }
    }

    fn transmute_slice_from_mut(val: &mut [T]) -> &mut [Self] {
        let len = val.len();
        let ptr = val.as_mut_ptr() as *mut Self;

        unsafe { core::slice::from_raw_parts_mut(ptr, len) }
    }
}

pub unsafe trait TransmutableInto<T: Sized>: Sized {
    fn transmute_into(val: Self) -> T;

    fn transmute_slice_into(val: &[Self]) -> &[T] {
        let len = val.len();
        let ptr = val.as_ptr() as *const T;

        unsafe { core::slice::from_raw_parts(ptr, len) }
    }

    fn transmute_slice_into_mut(val: &mut [Self]) -> &mut [T] {
        let len = val.len();
        let ptr = val.as_mut_ptr() as *mut T;

        unsafe { core::slice::from_raw_parts_mut(ptr, len) }
    }
}

unsafe impl<T: TransmutableFrom<F>, F: Sized> TransmutableInto<T> for F {
    fn transmute_into(val: Self) -> T {
        T::transmute_from(val)
    }
}

unsafe impl<T: Sized> TransmutableFrom<T> for T {
    fn transmute_from(val: T) -> Self {
        val
    }
}

pub trait SliceExt {
    type Item;

    /// Returns an iterator over `N` elements of the slice at a time, starting at the
    /// beginning of the slice.
    ///
    /// The chunks are array references and do not overlap. If `N` does not divide the
    /// length of the slice, then the last up to `N-1` elements will be omitted and can be
    /// retrieved from the [`remainder`](ArrayChunks::remainder) function of the iterator.
    ///
    /// Note: this function is designed to be equivalent to the currently unstable core::slice::array_chunks.
    fn uarray_chunks<const SIZE: usize>(&self) -> ArrayChunks<Self::Item, SIZE>;

    /// Returns an iterator over `N` elements of the slice at a time, starting at the
    /// beginning of the slice.
    ///
    /// The chunks are mutable array references and do not overlap. If `N` does not divide
    /// the length of the slice, then the last up to `N-1` elements will be omitted and
    /// can be retrieved from the [`into_remainder`](ArrayChunksMut::into_remainder) function of the iterator.
    ///
    /// Note: this function is designed to be equivalent to the currently unstable core::slice::array_chunks_mut.
    fn uarray_chunks_mut<const SIZE: usize>(&mut self) -> ArrayChunksMut<Self::Item, SIZE>;
}

impl<T> SliceExt for [T] {
    type Item = T;

    fn uarray_chunks<const SIZE: usize>(&self) -> ArrayChunks<Self::Item, SIZE> {
        ArrayChunks { inner: &self }
    }

    fn uarray_chunks_mut<const SIZE: usize>(&mut self) -> ArrayChunksMut<Self::Item, SIZE> {
        let len = self.len();
        let start = self.as_mut_ptr();

        ArrayChunksMut {
            start,
            len,
            phantomdata: PhantomData,
        }
    }
}

pub struct ArrayChunks<'a, T: 'a, const SIZE: usize> {
    inner: &'a [T],
}

impl<'a, T: 'a, const SIZE: usize> ArrayChunks<'a, T, SIZE> {
    pub fn remainder(&self) -> &'a [T] {
        self.inner
    }
}

impl<'a, T: 'a, const SIZE: usize> Iterator for ArrayChunks<'a, T, SIZE> {
    type Item = &'a [T; SIZE];

    fn next(&mut self) -> Option<Self::Item> {
        if self.inner.len() < SIZE {
            return None;
        }

        let (arr, rem) = self.inner.split_at(SIZE);
        self.inner = rem;

        let item = unsafe { <&'a [T; SIZE]>::try_from(arr).unwrap_unchecked() };

        Some(item)
    }
}

// I had to implement this using raw pointers because the compiler was giving weird lifetime errors
pub struct ArrayChunksMut<'a, T: 'a, const SIZE: usize> {
    start: *mut T,
    len: usize,

    phantomdata: PhantomData<&'a mut ()>,
}

impl<'a, T: 'a, const SIZE: usize> ArrayChunksMut<'a, T, SIZE> {
    pub fn into_remainder(self) -> &'a mut [T] {
        unsafe { core::slice::from_raw_parts_mut(self.start, self.len) }
    }
}

impl<'a, T: 'a, const SIZE: usize> Iterator for ArrayChunksMut<'a, T, SIZE> {
    type Item = &'a mut [T; SIZE];

    fn next(&mut self) -> Option<Self::Item> {
        if self.len < SIZE {
            return None;
        }

        unsafe {
            let arr: *mut [T; SIZE] = self.start.cast();

            self.len -= SIZE;
            self.start = self.start.offset(SIZE as isize);

            Some(&mut *arr)
        }
    }
}

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

    #[test]
    fn array_chunks() {
        let arr = [1, 2, 3, 5, 7, 11];
        let mut chunks = arr.uarray_chunks();

        assert_eq!(chunks.next(), Some(&[1, 2]));
        assert_eq!(chunks.next(), Some(&[3, 5]));
        assert_eq!(chunks.next(), Some(&[7, 11]));
        assert_eq!(chunks.next(), None);

        let arr = [1, 2, 3, 5, 7, 11, 13];
        let mut chunks = arr.uarray_chunks();

        assert_eq!(chunks.next(), Some(&[1, 2]));
        assert_eq!(chunks.next(), Some(&[3, 5]));
        assert_eq!(chunks.next(), Some(&[7, 11]));
        assert_eq!(chunks.next(), None);
    }

    #[test]
    fn array_chunks_mut() {
        let mut arr = [1, 2, 3, 5, 7, 11];
        let mut chunks = arr.uarray_chunks_mut();

        assert_eq!(chunks.next(), Some(&mut [1, 2]));
        assert_eq!(chunks.next(), Some(&mut [3, 5]));
        assert_eq!(chunks.next(), Some(&mut [7, 11]));
        assert_eq!(chunks.next(), None);

        let mut arr = [1, 2, 3, 5, 7, 11, 13];
        let mut chunks = arr.uarray_chunks_mut();

        assert_eq!(chunks.next(), Some(&mut [1, 2]));
        assert_eq!(chunks.next(), Some(&mut [3, 5]));
        assert_eq!(chunks.next(), Some(&mut [7, 11]));
        assert_eq!(chunks.next(), None);
    }
}