use core::marker::PhantomData;
use std::ptr::NonNull;

use crate::machine::heap::AllocError;
use crate::raw_block::*;
use crate::types::*;

use std::mem;
use std::ops::{Index, IndexMut};
use std::ptr;

impl RawBlockTraits for Stack {
    #[inline]
    fn init_size() -> usize {
        10 * 1024 * 1024
    }

    #[inline]
    fn align() -> usize {
        mem::align_of::<OrFrame>()
            .max(mem::align_of::<AndFrame>())
            .max(mem::align_of::<HeapCellValue>())
    }
}

#[inline(always)]
pub const fn prelude_size<Prelude>() -> usize {
    mem::size_of::<Prelude>()
}

#[derive(Debug)]
pub struct Stack {
    buf: RawBlock<Stack>,
    _marker: PhantomData<HeapCellValue>,
}

#[derive(Debug)]
pub(crate) struct AndFramePrelude {
    pub(crate) num_cells: usize,
    pub(crate) e: usize,
    pub(crate) cp: usize,
}

#[derive(Debug)]
pub(crate) struct AndFrame {
    pub(crate) prelude: AndFramePrelude,
}

impl AndFrame {
    pub(crate) fn size_of(num_cells: usize) -> usize {
        prelude_size::<AndFramePrelude>() + num_cells * mem::size_of::<HeapCellValue>()
    }
}

impl Index<usize> for AndFrame {
    type Output = HeapCellValue;

    fn index(&self, index: usize) -> &Self::Output {
        let prelude_offset = prelude_size::<AndFramePrelude>();
        let index_offset = (index - 1) * mem::size_of::<HeapCellValue>();

        unsafe {
            let ptr = self as *const crate::machine::stack::AndFrame as *const u8;

            // This address falls outside the provenance for self, therefore we have to get it
            // from exposed provenance.
            &*std::ptr::with_exposed_provenance(ptr.addr() + prelude_offset + index_offset)
        }
    }
}

impl IndexMut<usize> for AndFrame {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let prelude_offset = prelude_size::<AndFramePrelude>();
        let index_offset = (index - 1) * mem::size_of::<HeapCellValue>();

        unsafe {
            let ptr = self as *mut crate::machine::stack::AndFrame as *mut u8;

            // This address falls outside the provenance for self, therefore we have to get it
            // from exposed provenance.
            &mut *std::ptr::with_exposed_provenance_mut(ptr.addr() + prelude_offset + index_offset)
        }
    }
}

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

    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        unsafe {
            let ptr = self.buf.get_unchecked(index);
            &*ptr.cast::<HeapCellValue>()
        }
    }
}

impl IndexMut<usize> for Stack {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        unsafe {
            let ptr = self.buf.get_unchecked(index);
            &mut *ptr.cast_mut().cast::<HeapCellValue>()
        }
    }
}

#[derive(Debug)]
pub(crate) struct OrFramePrelude {
    pub(crate) num_cells: usize,
    pub(crate) e: usize,
    pub(crate) cp: usize,
    pub(crate) b: usize,
    pub(crate) bp: usize,
    pub(crate) boip: u32,
    pub(crate) biip: u32,
    pub(crate) tr: usize,
    pub(crate) h: usize,
    pub(crate) b0: usize,
    pub(crate) attr_var_queue_len: usize,
}

#[derive(Debug)]
pub(crate) struct OrFrame {
    pub(crate) prelude: OrFramePrelude,
}

impl Index<usize> for OrFrame {
    type Output = HeapCellValue;

    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        let prelude_offset = prelude_size::<OrFramePrelude>();
        let index_offset = index * mem::size_of::<HeapCellValue>();

        unsafe {
            let ptr = self as *const crate::machine::stack::OrFrame as *const u8;

            // This address falls outside the provenance for self, therefore we have to get it
            // from exposed provenance.
            &*std::ptr::with_exposed_provenance(ptr.addr() + prelude_offset + index_offset)
        }
    }
}

impl IndexMut<usize> for OrFrame {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let prelude_offset = prelude_size::<OrFramePrelude>();
        let index_offset = index * mem::size_of::<HeapCellValue>();

        unsafe {
            let ptr = self as *mut crate::machine::stack::OrFrame as *mut u8;

            // This address falls outside the provenance for self, therefore we have to get it
            // from exposed provenance.
            &mut *std::ptr::with_exposed_provenance_mut(ptr.addr() + prelude_offset + index_offset)
        }
    }
}

impl OrFrame {
    pub(crate) fn size_of(num_cells: usize) -> usize {
        prelude_size::<OrFramePrelude>() + num_cells * mem::size_of::<HeapCellValue>()
    }
}

impl Stack {
    pub(crate) fn new() -> Result<Self, AllocError> {
        Ok(Stack {
            buf: RawBlock::new()?,
            _marker: PhantomData,
        })
    }

    #[inline(always)]
    unsafe fn alloc(&mut self, frame_size: usize) -> Result<NonNull<u8>, AllocError> {
            loop {
                unsafe {
                let ptr = self.buf.alloc(frame_size);
                if let Some(ptr) = NonNull::new(ptr) {
                    return Ok(ptr);
                }
                self.buf.grow()?;
            }
        }
    }

    pub(crate) fn allocate_and_frame(&mut self, num_cells: usize) -> Result<usize, AllocError> {
        let frame_size = AndFrame::size_of(num_cells);

        unsafe {
            let e = self.buf.used_bytes();
            let new_ptr = self.alloc(frame_size)?;
            let mut offset = prelude_size::<AndFramePrelude>();

            for idx in 0..num_cells {
                let cell_ptr = new_ptr.add(offset).cast::<HeapCellValue>();
                ptr::write(cell_ptr.as_ptr(), stack_loc_as_cell!(AndFrame, e, idx + 1));

                // Because in the Index and IndexMut implementations we need to get this from
                // exposed provenance, we need to expose the provenance here, even though we don't
                // actually use the value for anything. This is a reminder that `expose_provenance`
                // isn't just a cast from a pointer to an integer but has actual side effects.
                cell_ptr.expose_provenance();

                offset += mem::size_of::<HeapCellValue>();
            }

            let and_frame = self.index_and_frame_mut(e);
            and_frame.prelude.num_cells = num_cells;

            Ok(e)
        }
    }

    pub(crate) fn top(&self) -> usize {
        self.buf.used_bytes()
    }

    pub(crate) fn allocate_or_frame(&mut self, num_cells: usize) -> Result<usize, AllocError> {
        let frame_size = OrFrame::size_of(num_cells);

        unsafe {
            let b = self.buf.used_bytes();
            let new_ptr = self.alloc(frame_size)?;
            let mut offset = prelude_size::<OrFramePrelude>();

            for idx in 0..num_cells {
                let cell_ptr = new_ptr.byte_add(offset).cast::<HeapCellValue>();
                ptr::write(cell_ptr.as_ptr(), stack_loc_as_cell!(OrFrame, b, idx));

                // Because in the Index and IndexMut implementations we need to get this from
                // exposed provenance, we need to expose the provenance here, even though we don't
                // actually use the value for anything. This is a reminder that `expose_provenance`
                // isn't just a cast from a pointer to an integer but has actual side effects.
                cell_ptr.expose_provenance();

                offset += mem::size_of::<HeapCellValue>();
            }

            let or_frame = self.index_or_frame_mut(b);
            or_frame.prelude.num_cells = num_cells;

            Ok(b)
        }
    }

    fn get_raw(&self, index: usize) -> *const u8 {
        debug_assert!(index < self.buf.used_bytes());

        unsafe { self.buf.get_unchecked(index) }
    }

    #[inline(always)]
    pub(crate) fn index_and_frame(&self, e: usize) -> &AndFrame {
        let ptr = self.get_raw(e);

        unsafe { &*ptr.cast::<AndFrame>() }
    }

    #[inline(always)]
    pub(crate) fn index_and_frame_mut(&mut self, e: usize) -> &mut AndFrame {
        let ptr = self.get_raw(e);

        unsafe { &mut *ptr.cast_mut().cast::<AndFrame>() }
    }

    #[inline(always)]
    pub(crate) fn index_or_frame(&self, b: usize) -> &OrFrame {
        let ptr = self.get_raw(b);

        unsafe { &*ptr.cast::<OrFrame>() }
    }

    #[inline(always)]
    pub(crate) fn index_or_frame_mut(&mut self, b: usize) -> &mut OrFrame {
        let ptr = self.get_raw(b);

        unsafe { &mut *ptr.cast_mut().cast::<OrFrame>() }
    }

    /// # Safety
    ///
    /// The stack must contain a valid OrFrame at [`self.top()`](Self::top),
    /// which can only be achieved by allocating it in the first place and later truncating the stack.
    ///
    /// No allocation must have been done since the last call to [`truncate()`](Self::truncate).
    #[inline(always)]
    pub(crate) unsafe fn index_dangling_or_frame(&self) -> &OrFrame {
        unsafe {
            let ptr = self.buf.get_unchecked(self.top());
            &*ptr.cast::<OrFrame>()
        }
    }

    #[inline(always)]
    pub(crate) fn truncate(&mut self, b: usize) {
        self.buf.shift_back(b);
    }
}

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

    use crate::machine::mock_wam::*;

    #[test]
    fn stack_tests() {
        let mut wam = MockWAM::new();

        let e = wam.machine_st.stack.allocate_and_frame(10).unwrap(); // create an AND frame!
        let and_frame = wam.machine_st.stack.index_and_frame_mut(e);

        assert_eq!(
            e,
            0 // 10 * mem::size_of::<HeapCellValue>() + prelude_size::<AndFrame>()
        );

        assert_eq!(and_frame.prelude.num_cells, 10);

        for idx in 0..10 {
            assert_eq!(and_frame[idx + 1], stack_loc_as_cell!(AndFrame, e, idx + 1));
        }

        and_frame[5] = empty_list_as_cell!();

        assert_eq!(and_frame[5], empty_list_as_cell!());

        let b = wam.machine_st.stack.allocate_or_frame(5).unwrap();

        let or_frame = wam.machine_st.stack.index_or_frame_mut(b);

        for idx in 0..5 {
            assert_eq!(or_frame[idx], stack_loc_as_cell!(OrFrame, b, idx));
        }

        let next_e = wam.machine_st.stack.allocate_and_frame(9).unwrap(); // create an AND frame!
        let and_frame = wam.machine_st.stack.index_and_frame_mut(next_e);

        for idx in 0..9 {
            assert_eq!(
                and_frame[idx + 1],
                stack_loc_as_cell!(AndFrame, next_e, idx + 1)
            );
        }

        let and_frame = wam.machine_st.stack.index_and_frame(e);
        assert_eq!(and_frame[5], empty_list_as_cell!());

        assert_eq!(
            wam.machine_st.stack[stack_loc!(AndFrame, e, 5)],
            empty_list_as_cell!()
        );
    }
}