use std::{
    mem::MaybeUninit,
    ops::{Bound, Index, RangeBounds},
};

pub struct Stack<T> {
    data: Box<[MaybeUninit<T>]>,
    pointer: usize,
}

impl<T> Stack<T> {
    pub fn new(limit: usize) -> Self {
        assert_ne!(limit, 0);
        Self {
            data: Box::new_uninit_slice(limit),
            pointer: 0,
        }
    }

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

    pub fn checked_slice<I: RangeBounds<usize>>(&self, range: I) -> Option<&[T]> {
        let start = match range.start_bound() {
            Bound::Unbounded => 0,
            Bound::Included(start) | Bound::Excluded(start) => *start,
        };
        let end = match range.end_bound() {
            Bound::Unbounded => self.pointer,
            Bound::Excluded(end) => *end,
            Bound::Included(end) => *end + 1,
        };

        if start <= self.pointer && end <= self.pointer {
            Some(unsafe { self.data[start..end].assume_init_ref() })
        } else {
            None
        }
    }

    pub fn checked_last(&self, n: usize) -> Option<&[T]> {
        if n < self.pointer {
            Some(unsafe { self.data[self.pointer - n..self.pointer].assume_init_ref() })
        } else {
            None
        }
    }

    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        if index < self.pointer {
            Some(unsafe { self.data[index].assume_init_mut() })
        } else {
            None
        }
    }

    pub fn nth(&self, n: usize) -> Option<&T> {
        if n >= self.pointer {
            None
        } else {
            Some(unsafe { self.data[self.pointer - n - 1].assume_init_ref() })
        }
    }

    pub fn get(&self, index: usize) -> Option<&T> {
        if index < self.pointer {
            Some(unsafe { self.data[index].assume_init_ref() })
        } else {
            None
        }
    }

    pub fn head(&self) -> Option<&T> {
        if self.pointer > 0 {
            Some(unsafe { self.data[self.pointer - 1].assume_init_ref() })
        } else {
            None
        }
    }

    pub fn head_mut(&mut self) -> Option<&mut T> {
        if self.pointer > 0 {
            Some(unsafe { self.data[self.pointer - 1].assume_init_mut() })
        } else {
            None
        }
    }

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

    pub fn set_pointer(&mut self, pointer: usize) {
        if pointer < self.pointer {
            // Call drop for values between
            for slot in &mut self.data[pointer..self.pointer] {
                unsafe { slot.assume_init_drop() };
            }
        } else if pointer > self.pointer {
            panic!("Cannot set stack pointer higher than it currently is");
        }

        self.pointer = pointer;
    }

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

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

impl<T> Index<usize> for Stack<T> {
    type Output = T;

    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        self.get(index).expect("stack index out of bounds")
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicBool, Ordering};

    use crate::vm::stack::Stack;

    #[derive(Debug)]
    struct DropTracker(&'static AtomicBool);

    impl Drop for DropTracker {
        fn drop(&mut self) {
            self.0.store(true, Ordering::Release);
        }
    }

    #[test]
    fn test_checked_slice() {
        let mut s = Stack::<u32>::new(32);
        s.push(1).unwrap();
        s.push(2).unwrap();
        s.push(3).unwrap();
        s.push(4).unwrap();
        s.push(5).unwrap();

        // unbound..unbound
        assert_eq!(s.checked_slice(..).unwrap(), &[1, 2, 3, 4, 5]);
        // unbound..included
        assert_eq!(s.checked_slice(..=3).unwrap(), &[1, 2, 3, 4]);
        // unbound..excluded
        assert_eq!(s.checked_slice(..3).unwrap(), &[1, 2, 3]);

        // included..unbound
        assert_eq!(s.checked_slice(2..).unwrap(), &[3, 4, 5]);
        // included..included
        assert_eq!(s.checked_slice(2..=3).unwrap(), &[3, 4]);
        // included..excluded
        assert_eq!(s.checked_slice(2..3).unwrap(), &[3]);

        // Invalid
        assert!(s.checked_slice(..10).is_none());
        assert!(s.checked_slice(10..).is_none());
        assert!(s.checked_slice(10..9).is_none());
    }

    #[test]
    fn test_underflow() {
        let mut s = Stack::<u32>::new(128);
        assert!(s.pop().is_none());
    }

    #[test]
    fn test_overflow() {
        let mut s = Stack::<u32>::new(4);
        for i in 0..4 {
            s.push(i).unwrap();
        }
        assert_eq!(s.push(1234).unwrap_err(), 1234);
    }

    #[test]
    fn test_no_drop_on_pop() {
        static DROP0: AtomicBool = AtomicBool::new(false);
        let mut s = Stack::<DropTracker>::new(4);
        s.push(DropTracker(&DROP0)).unwrap();
        assert!(!DROP0.load(Ordering::Acquire));
        let v = s.pop().unwrap();
        assert!(!DROP0.load(Ordering::Acquire));
        drop(v);
        assert!(DROP0.load(Ordering::Acquire));
    }

    #[test]
    fn test_drop_on_set_pointer_lower() {
        static DROP0: AtomicBool = AtomicBool::new(false);
        static DROP1: AtomicBool = AtomicBool::new(false);

        let mut s = Stack::<DropTracker>::new(4);
        s.push(DropTracker(&DROP0)).unwrap();
        s.push(DropTracker(&DROP1)).unwrap();
        s.set_pointer(1);
        assert!(!DROP0.load(Ordering::Acquire));
        assert!(DROP1.load(Ordering::Acquire));
    }

    #[test]
    fn test_stack_head() {
        let mut s = Stack::<u32>::new(4);
        assert!(s.head().is_none());
        s.push(1234).unwrap();
        assert_eq!(s.head().unwrap(), &1234);
    }
}
