// 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: u64 = 8192;

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

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, Layout, handle_alloc_error};
            let layout = Layout::new::<[u8; PAGE_SIZE as usize]>();
            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 + u16::from(S::BYTES) <= (PAGE_SIZE as u16 - 1) {
            let mut value = MaybeUninit::<Value>::zeroed();
            unsafe {
                core::ptr::copy_nonoverlapping(
                    self.pages.get(&page)?.as_ptr().add(usize::from(offset)),
                    value.as_mut_ptr().cast(),
                    S::BYTES.into(),
                );
                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 + u16::from(S::BYTES) <= (PAGE_SIZE as u16 - 1) {
            unsafe {
                core::ptr::copy_nonoverlapping(
                    (&value as *const Value).cast::<u8>(),
                    self.pages
                        .get_mut(&page)
                        .ok_or(())?
                        .as_mut_ptr()
                        .add(usize::from(offset)),
                    S::BYTES.into(),
                )
            };
            Ok(())
        } else {
            Err(())
        }
    }
}

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

macro_rules! size_markers {
    ($($name:ident = $size:expr),* $(,)?) => {
        pub mod ma_size {
            pub unsafe trait MemAccessSize {
                const BYTES: u8;
            }

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

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