use core::mem::MaybeUninit;

#[derive(Debug)]
pub struct StaticVec<T, const N: usize> {
    data: [MaybeUninit<T>; N],
    len: usize,
}

impl<T: Clone, const N: usize> Clone for StaticVec<T, N> {
    fn clone(&self) -> Self {
        let mut new = Self::new();
        self.as_slice().iter().for_each(|item| {
            let _ = new.push(item.clone());
        });
        new
    }
}

impl<T, const N: usize> StaticVec<T, N> {
    pub const fn new() -> Self {
        Self {
            data: [const { MaybeUninit::uninit() }; N],
            len: 0,
        }
    }

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

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

    #[allow(dead_code)]
    pub const fn capacity(&self) -> usize {
        N
    }

    #[allow(dead_code)]
    pub const fn is_full(&self) -> bool {
        self.len == N
    }

    pub fn push(&mut self, val: T) -> Result<(), T> {
        if self.len < N {
            self.data[self.len] = MaybeUninit::new(val);
            self.len += 1;
            Ok(())
        } else {
            Err(val)
        }
    }

    pub fn pop(&mut self) -> Option<T> {
        if self.len > 0 {
            self.len -= 1;
            Some(unsafe { self.data[self.len].assume_init_read() })
        } else {
            None
        }
    }

    pub fn as_slice(&self) -> &[T] {
        unsafe { core::slice::from_raw_parts(self.data.as_ptr() as *const T, self.len) }
    }

    #[allow(dead_code)]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        unsafe { core::slice::from_raw_parts_mut(self.data.as_mut_ptr() as *mut T, self.len) }
    }

    pub fn get(&self, index: usize) -> Option<&T> {
        self.as_slice().get(index)
    }

    #[allow(dead_code)]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        self.as_mut_slice().get_mut(index)
    }

    pub fn clear(&mut self) {
        (0..self.len).rev().for_each(|i| unsafe {
            core::ptr::drop_in_place(self.data[i].as_mut_ptr());
        });
        self.len = 0;
    }

    pub fn iter(&self) -> core::slice::Iter<'_, T> {
        self.as_slice().iter()
    }

    #[allow(dead_code)]
    pub fn iter_mut(&mut self) -> core::slice::IterMut<'_, T> {
        self.as_mut_slice().iter_mut()
    }
}

impl<T, const N: usize> Drop for StaticVec<T, N> {
    fn drop(&mut self) {
        self.clear();
    }
}

impl<T, const N: usize> StaticVec<T, N> {
    pub fn swap_remove(&mut self, index: usize) -> T {
        assert!(index < self.len, "swap_remove index out of bounds");
        self.len -= 1;
        let removed = unsafe { self.data[index].assume_init_read() };
        if index != self.len {
            self.data[index] = unsafe { MaybeUninit::new(self.data[self.len].assume_init_read()) };
        }
        removed
    }
}

#[allow(dead_code)]
impl<T: Copy, const N: usize> StaticVec<T, N> {
    pub fn retain(&mut self, pred: impl Fn(&T) -> bool) {
        let write = (0..self.len).fold(0usize, |write, read| {
            if pred(&self.as_slice()[read]) {
                if write != read {
                    self.data[write] = self.data[read];
                }
                write + 1
            } else {
                write
            }
        });
        self.len = write;
    }
}

#[cfg(kani)]
mod kani_proofs {
    use super::*;

    #[kani::proof]
    #[kani::unwind(6)]
    fn push_pop_inverse() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();
        let val: u32 = kani::any();
        if v.push(val).is_ok() {
            let popped = v.pop();
            assert!(popped == Some(val));
        }
    }

    #[kani::proof]
    #[kani::unwind(6)]
    fn len_invariant_after_ops() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();
        let op_count: usize = kani::any();
        kani::assume(op_count <= 4);

        (0..op_count).for_each(|_| {
            let val: u32 = kani::any();
            let _ = v.push(val);
        });

        assert!(v.len() <= 4);
        assert!(v.as_slice().len() == v.len());
    }

    #[kani::proof]
    #[kani::unwind(6)]
    fn push_full_returns_err() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();
        (0..4u32).for_each(|i| {
            let _ = v.push(i);
        });
        assert!(v.is_full());
        let result = v.push(99);
        assert!(result.is_err());
        assert!(result.err() == Some(99));
    }

    #[kani::proof]
    #[kani::unwind(6)]
    fn pop_empty_returns_none() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();
        assert!(v.pop().is_none());
    }

    #[kani::proof]
    #[kani::unwind(6)]
    fn swap_remove_preserves_len() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();
        let count: usize = kani::any();
        kani::assume(count > 0 && count <= 4);

        (0..count).for_each(|i| {
            let _ = v.push(i as u32);
        });

        let idx: usize = kani::any();
        kani::assume(idx < v.len());

        let old_len = v.len();
        let _ = v.swap_remove(idx);
        assert!(v.len() == old_len - 1);
    }

    #[kani::proof]
    #[kani::unwind(6)]
    fn push_pop_sequence_preserves_invariant() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();

        let ops: [u8; 4] = kani::any();
        ops.iter().for_each(|&op| {
            match op % 3 {
                0 => {
                    let val: u32 = kani::any();
                    let _ = v.push(val);
                }
                1 => {
                    let _ = v.pop();
                }
                _ => {}
            }
            assert!(v.len() <= 4);
            assert!(v.as_slice().len() == v.len());
        });
    }

    #[kani::proof]
    #[kani::unwind(6)]
    fn get_in_bounds_succeeds() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();
        let count: usize = kani::any();
        kani::assume(count > 0 && count <= 4);

        (0..count).for_each(|i| {
            let _ = v.push(i as u32);
        });

        let idx: usize = kani::any();
        kani::assume(idx < count);
        assert!(v.get(idx).is_some());

        let oob: usize = kani::any();
        kani::assume(oob >= count);
        assert!(v.get(oob).is_none());
    }

    #[kani::proof]
    #[kani::unwind(6)]
    fn retain_preserves_order_and_shrinks() {
        let mut v: StaticVec<u32, 4> = StaticVec::new();
        (0..4u32).for_each(|i| {
            let _ = v.push(i);
        });

        let threshold: u32 = kani::any();
        kani::assume(threshold <= 4);
        v.retain(|&x| x < threshold);

        assert!(v.len() <= 4);
        assert!(v.len() == threshold as usize);
    }
}