// HACK: This is temporary implementation so we can have memory instructions working

use {
    crate::vm::value::Value, alloc::boxed::Box, core::mem::MaybeUninit, hashbrown::HashMap,
    ma_size::MemAccessSize,
};

pub const PAGE_SIZE: usize = 8192;

#[derive(Clone, Debug, Default)]
pub struct Memory {
    pages: HashMap<u64, Box<[u8; PAGE_SIZE]>>,
}

impl Memory {
    // HACK: Just for allocation testing, will be removed when proper memory interfaces
    // implemented.
    pub fn insert_test_page(&mut self) {
        self.pages.insert(0, unsafe {
            use alloc::alloc::{alloc_zeroed, handle_alloc_error, Layout};
            let layout = Layout::new::<[u8; PAGE_SIZE]>();
            let ptr = alloc_zeroed(layout);
            if ptr.is_null() {
                handle_alloc_error(layout);
            }
            Box::from_raw(ptr.cast())
        });
    }

    pub fn load<S: MemAccessSize>(&self, addr: u64) -> Option<Value> {
        let (page, offset) = split_addr(addr);
        if offset + S::BYTES <= PAGE_SIZE - 1 {
            let mut value = MaybeUninit::<Value>::zeroed();
            unsafe {
                core::ptr::copy_nonoverlapping(
                    self.pages.get(&page)?.as_ptr().add(offset),
                    value.as_mut_ptr().cast(),
                    S::BYTES,
                );
                Some(value.assume_init())
            }
        } else {
            None
        }
    }

    pub fn store<S: MemAccessSize>(&mut self, addr: u64, value: Value) -> Result<(), ()> {
        let (page, offset) = split_addr(addr);
        if offset + S::BYTES <= PAGE_SIZE - 1 {
            unsafe {
                core::ptr::copy_nonoverlapping(
                    (&value as *const Value).cast::<u8>(),
                    self.pages
                        .get_mut(&page)
                        .ok_or(())?
                        .as_mut_ptr()
                        .add(offset),
                    S::BYTES,
                )
            };
            Ok(())
        } else {
            Err(())
        }
    }
}

#[inline]
pub const fn split_addr(addr: u64) -> (u64, usize) {
    (addr >> PAGE_SIZE.count_ones(), (addr as usize & PAGE_SIZE))
}

macro_rules! size_markers {
    ($($name:ident = $size:expr),* $(,)?) => {
        pub mod ma_size {
            /// # Safety
            /// Implementor has to assure that [`MemAccessSize::BYTES`] won't be larger than
            /// size of [`Value`]
            pub unsafe trait MemAccessSize {
                const BYTES: usize;
            }

            $(
                pub struct $name;
                unsafe impl MemAccessSize for $name {
                    const BYTES: usize = $size;
                }
            )*
        }
    };
}

size_markers! {
    Byte    = 1,
    Doublet = 2,
    Quadlet = 4,
    Octlet  = 8,
}